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

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

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

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

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

始める前に

依存関係のインストール

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

$ pip install pymilvus>=2.4.2

$ pip install llama-index-vector-stores-milvus

$ pip install llama-index

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

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

文書ができたので、インデックスを作成し、文書を挿入することができます。

Milvus Liteには pymilvus>=2.4.2 が必要です。

# 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)

データをクエリする

ドキュメントがインデックスに格納されたので、インデックスに対して質問することができます。インデックスはそれ自体に格納されたデータを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.