MilvusとLlamaIndexによる検索拡張生成(RAG)

Open In Colab GitHub Repository

このガイドでは、LlamaIndexとMilvusを使ったRAG(Retrieval-Augmented Generation)システムの構築方法を説明する。

RAGシステムは検索システムと生成モデルを組み合わせ、与えられたプロンプトに基づいて新しいテキストを生成する。システムはまずMilvusを使ってコーパスから関連文書を検索し、次に生成モデルを使って検索された文書に基づいて新しいテキストを生成する。

LlamaIndexは、カスタムデータソースを大規模言語モデル(LLM)に接続するためのシンプルで柔軟なデータフレームワークである。Milvusは世界で最も先進的なオープンソースのベクトル・データベースであり、埋め込み類似検索やAIアプリケーションのために構築されている。

このノートブックでは、MilvusVectorStoreの簡単なデモをお見せします。

始める前に

依存関係のインストール

このページのコードスニペットにはpymilvusとllamaindexの依存関係が必要です。以下のコマンドでインストールできます:

$ pip install pymilvus>=2.4.2 milvus-lite
$ pip install llama-index-vector-stores-milvus
$ pip install llama-index

Google Colabを使用している場合、インストールした依存関係を有効にするために、ランタイムを再起動する必要があるかもしれません。(画面上部の "Runtime "メニューをクリックし、ドロップダウンメニューから "Restart session "を選択してください)。

OpenAIのセットアップ

まず、openai api keyを追加します。これでchatgptにアクセスできるようになる。

import openai

openai.api_key = "sk-***********"

データの準備

以下のコマンドでサンプルデータをダウンロードできます:

! mkdir -p 'data/'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'

はじめに

データを生成する

最初の例として、paul_graham_essay.txt というファイルから文書を生成してみましょう。これはWhat I Worked On というタイトルの Paul Graham のエッセイです。ドキュメントを生成するにはSimpleDirectoryReaderを使います。

from llama_index.core import SimpleDirectoryReader

# load documents
documents = SimpleDirectoryReader(
    input_files=["./data/paul_graham_essay.txt"]
).load_data()

print("Document ID:", documents[0].doc_id)
Document ID: 95f25e4d-f270-4650-87ce-006d69d82033

データを横断するインデックスの作成

文書ができたので、インデックスを作成して文書を挿入することができます。インデックスにはMilvusVectorStoreを使います。MilvusVectorStoreはいくつかの引数をとります:

基本引数

  • uri (str, optional):接続先のURIは、MilvusやZillizクラウドサービスの場合は "https://address:port"、ローカルMilvusの場合は "path/to/local/milvus.db "となる。デフォルトは"./milvus_llamaindex.db "です。
  • token (str, optional):ログイン用トークン。rbacを使用しない場合は空、rbacを使用する場合は "username:password "となる。
  • collection_name (str, optional):データを保存するコレクション名。デフォルトは "llamalection"。
  • overwrite (bool, optional):既存の同じ名前のコレクションを上書きするかどうか。デフォルトはFalse。

doc idとtextを含むスカラーフィールド

  • doc_id_field (str, optional):コレクションのdoc_idフィールドの名前。デフォルトは DEFAULT_DOC_ID_KEY です。
  • text_key (str, optional):渡されたコレクションのどのキーテキストに格納されるか。独自のコレクションを持ってくるときに使用する。デフォルトは DEFAULT_TEXT_KEY です。
  • scalar_field_names (list, optional):コレクションスキーマに含まれる追加スカラーフィールドの名前。
  • scalar_field_types (list, optional):追加のスカラー・フィールドの型。

密フィールド

  • enable_dense (bool):密な埋め込みを有効または無効にするブーリアン・フラグ。デフォルトはTrue。
  • dim (int, optional):コレクションの埋め込みベクトルの次元。enable_sparse が False で新しいコレクションを作成するときに必要です。
  • embedding_field (str, optional):コレクションの密埋め込みフィールドの名前。デフォルトは DEFAULT_EMBEDDING_KEY です。
  • index_config (dict, optional):密埋め込みインデックスの構築に使用する設定。デフォルトはNone。
  • search_config (dict, optional):Milvus密インデックスを検索するための設定。これはindex_config で指定されたインデックスタイプと互換性がなければならないことに注意してください。デフォルトはなし。
  • similarity_metric (str, optional):密な埋め込みに使用する類似度メトリックで、現在は IP, COSINE, L2 をサポートしています。

スパースフィールド

  • enable_sparse (bool):スパース埋め込みを有効または無効にするブール値のフラグ。デフォルトはFalse。
  • sparse_embedding_field (str):スパース埋め込みフィールドの名前。デフォルトは DEFAULT_SPARSE_EMBEDDING_KEY です。
  • sparse_embedding_function (Union[BaseSparseEmbeddingFunction, BaseMilvusBuiltInFunction], optional):enable_sparse が True の場合、テキストをスパース埋め込みに変換するためにこのオブジェクトを提供する必要があります。Noneの場合、デフォルトのスパース埋め込み関数(BGEM3SparseEmbeddingFunction)が使用される。
  • sparse_index_config (dict, optional):スパース埋め込みインデックスの構築に使用される設定。デフォルトはNoneです。

ハイブリッドランカー

  • hybrid_ranker (str):ハイブリッド検索クエリで使用するランカーのタイプを指定します。現在は["RRFRanker", "WeightedRanker"]のみサポート。デフォルトは "RRFRanker"。

  • hybrid_ranker_params (dict, optional):ハイブリッドランカーの設定パラメータ。この辞書の構造は使用される特定のランカーに依存する:

    • RRFRanker "の場合、以下を含むべきである:
      • "k"(int):k」(int):RRF(Reciprocal Rank Fusion)で使用されるパラメータ。この値はRRFアルゴリズムの一部としてランクスコアを計算するために使用され、複数のランキング戦略を1つのスコアにまとめ、検索の関連性を向上させます。
    • WeightedRanker "には、次のようなものが期待されます:
      • "weights"(floatのリスト):2つの重みのリスト:
        1. 密な埋め込みコンポーネントの重み。
        2. これらの重みは、ハイブリッド検索プロセスにおいて、埋め込み要素の密な成分と疎な成分の重要度を調整するために使用されます。 デフォルトは空の辞書で、ランカーは事前に定義されたデフォルト設定で動作することを意味します。

その他

  • collection_properties (dict, optional):TTL (Time-To-Live)やMMAP (memory mapping)などのコレクション・プロパティ。デフォルトはNone。以下を含むことができる:
    • "collection.ttl.seconds"(int):このプロパティが設定されると、現在のコレクション内のデータは指定された時間で期限切れになる。コレクション内の期限切れデータはクリーンアップされ、検索やクエリに関与しない。
    • "mmap.enabled" (bool):コレクション・レベルでメモリ・マップド・ストレージを有効にするかどうか。
  • index_management (IndexManagement):使用するインデックス管理ストラテジを指定します。デフォルトは "create_if_not_exists"。
  • batch_size (int):Milvusへのデータ挿入時に1バッチで処理する文書数を設定します。デフォルトはDEFAULT_BATCH_SIZEです。
  • consistency_level (str, optional):新しく作成されたコレクションにどの一貫性レベルを使用するか。デフォルトは "Session "です。
# Create an index over the documents
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

MilvusVectorStore のパラメータについて:

  • uri をローカルファイル、例えば./milvus.db に設定するのが最も便利な方法である。
  • データ規模が大きい場合は、dockerやkubernetes上に、よりパフォーマンスの高いMilvusサーバを構築することができます。このセットアップでは、サーバの uri、例えばhttp://localhost:19530uri として使用してください。
  • MilvusのフルマネージドクラウドサービスであるZilliz Cloudを利用する場合は、Zilliz CloudのPublic EndpointとApi keyに対応するuritoken を調整してください。

データをクエリする

ドキュメントがインデックスに格納されたので、インデックスに対して質問をすることができます。インデックスはそれ自体に格納されたデータをchatgptの知識ベースとして使用する。

query_engine = index.as_query_engine()
res = query_engine.query("What did the author learn?")
print(res)
The author learned that philosophy courses in college were boring to him, leading him to switch his focus to studying AI.
res = query_engine.query("What challenges did the disease pose for the author?")
print(res)
The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in her losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.

次のテストでは、上書きによって前のデータが削除されることを示しています。

from llama_index.core import Document


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    [Document(text="The number that is being searched for is ten.")],
    storage_context,
)
query_engine = index.as_query_engine()
res = query_engine.query("Who is the author?")
print(res)
The author is the individual who created the context information.

次のテストは、既に存在するインデックスにデータを追加することを示しています。

del index, vector_store, storage_context, query_engine

vector_store = MilvusVectorStore(uri="./milvus_demo.db", overwrite=False)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
query_engine = index.as_query_engine()
res = query_engine.query("What is the number?")
print(res)
The number is ten.
res = query_engine.query("Who is the author?")
print(res)
Paul Graham

メタデータのフィルタリング

特定のソースをフィルタリングして結果を生成することができます。次の例では、ディレクトリからすべての文書を読み込み、その後メタデータに基づいてそれらをフィルタリングしています。

from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

# Load all the two documents loaded before
documents_all = SimpleDirectoryReader("./data/").load_data()

vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents_all, storage_context)

uber_2021.pdf というファイルから文書だけを取り出したいとします。

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="uber_2021.pdf")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)
The disease posed challenges related to the adverse impact on the business and operations, including reduced demand for Mobility offerings globally, affecting travel behavior and demand. Additionally, the pandemic led to driver supply constraints, impacted by concerns regarding COVID-19, with uncertainties about when supply levels would return to normal. The rise of the Omicron variant further affected travel, resulting in advisories and restrictions that could adversely impact both driver supply and consumer demand for Mobility offerings.

こ こ で フ ァ イ ルpaul_graham_essay.txt か ら 文書を取得す る と 、 別の結果が得 ら れます。

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="paul_graham_essay.txt")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)
The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in his mother losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.