MilvusとEmbedAnythingでRAGを構築する

EmbedAnythingは、テキスト、PDF、画像、音声などをサポートする、Rustで構築された高速で軽量な埋め込みパイプラインです。

このチュートリアルでは、EmbedAnythingとMilvusを使用したRAG（Retrieval-Augmented Generation）パイプラインの構築方法を紹介します。EmbedAnythingは、特定のデータベースと密に結合するのではなく、プラグイン可能なアダプタシステムを使用しています。アダプタは、エンベッディングがどのようにフォーマットされ、インデックス付けされ、ターゲットベクターストアに格納されるかを定義するラッパーとして機能します。

EmbedAnythingとMilvusアダプタを組み合わせることで、わずか数行のコードで、多様なファイルタイプからエンベッディングを生成し、効率的にMilvusに格納することができます。

⚠️ 注：EmbedAnythingのアダプタはMilvusへの挿入を処理しますが、検索はサポートしていません。完全なRAGパイプラインを構築するには、別途MilvusClientをインスタンス化し、アプリケーションの一部として検索ロジック（例：ベクトルに対する類似検索）を実装する必要があります。

準備

依存関係と環境

$ pip install -qU pymilvus milvus-lite openai embed_anything

リポジトリのクローンとアダプタのロード

次に、EmbedAnythingリポジトリをクローンし、examples/adapters ディレクトリを Python のパスに追加します。このディレクトリにはEmbedAnythingがMilvusと通信してベクトルを挿入するためのカスタムMilvusアダプタ実装が格納されています。

import sys

# Clone the EmbedAnything repository if not already cloned
![ -d "EmbedAnything" ] || git clone https://github.com/StarlightSearch/EmbedAnything.git

# Add the `examples/adapters` directory to the Python path
sys.path.append("EmbedAnything/examples/adapters")
print("✅ EmbedAnything cloned and adapter path added.")

✅ EmbedAnything cloned and adapter path added.

このRAGパイプラインではOpenAIをLLMとして使用します。環境変数としてapi key OPENAI_API_KEY を用意してください。

import os
from openai import OpenAI

os.environ["OPENAI_API_KEY"] = "sk-***********"

openai_client = OpenAI()

RAGのビルド

Milvusの初期化

ファイルを埋め込む前に、Milvusとやりとりする2つのコンポーネントを準備する必要がある：

1. MilvusVectorAdapter - これはEmbedAnything用のMilvusアダプタで、ベクトルインジェスト（エンベッディングの挿入とインデックスの作成）のみに使用します。現在のところ、検索操作はサポートしていません。
2. MilvusClient - これはpymilvus からの公式クライアントであり、ベクター検索、フィルタリング、コレクション管理など、Milvus の全機能にアクセスすることができます。

混乱を避けるために

• MilvusVectorAdapter はベクターを保存するための "書き込み専用 "ツールであるとお考えください。
• MilvusClient は、実際にクエリーを実行し、RAGのドキュメントを検索するための "読み取り・検索 "エンジンであると考えてください。

import embed_anything
from embed_anything import (
    WhichModel,
    EmbeddingModel,
)
from milvus_db import MilvusVectorAdapter
from pymilvus import MilvusClient

# Official Milvus client for full operations
milvus_client = MilvusClient(uri="./milvus.db", token="")

# EmbedAnything adapter for pushing embeddings into Milvus
index_name = "embed_anything_milvus_collection"
milvus_adapter = MilvusVectorAdapter(
    uri="./milvus.db", token="", collection_name=index_name
)

# Delete existing collection if it exists
if milvus_client.has_collection(index_name):
    milvus_client.drop_collection(index_name)

# Create a new collection with dimension matching the embedding model later used
milvus_adapter.create_index(dimension=384)

Ok - Milvus DB connection established.
Collection 'embed_anything_milvus_collection' created with index.

埋め込みモデルの初期化とPDFドキュメントの埋め込み

埋め込みモデルを初期化します。Sentence-transformersライブラリのall-MiniLM-L12-v2 model 。これは軽量でありながら、テキスト埋め込みを生成するための強力なモデルです。これは384次元の埋め込みを生成するので、Milvusコレクションの次元が384に設定されているのと一致します。このアラインメントは非常に重要で、Milvusに保存されたベクトル次元とモデルによって生成されたベクトル次元の互換性を保証します。

EmbedAnythingはさらに多くの埋め込みモデルをサポートしています。詳しくは公式ドキュメントを参照してください。

# Initialize the embedding model
model = EmbeddingModel.from_pretrained_hf(
    WhichModel.Bert, model_id="sentence-transformers/all-MiniLM-L12-v2"
)

それでは、PDFファイルを埋め込んでみましょう。EmbedAnythingは、PDF（およびその他多くの）ドキュメントを簡単に処理し、Milvusに直接埋め込むことができます。

# Embed a PDF file
data = embed_anything.embed_file(
    "./pdf_files/WhatisMilvus.pdf",
    embedder=model,
    adapter=milvus_adapter,
)

Converted 12 embeddings for insertion.
Successfully inserted 12 embeddings.

レスポンスの取得と生成

繰り返しますが、現在のところEmbedAnythingのMilvusVectorAdapter は、ベクトルの取り込みとインデックス作成のみのための軽量な抽象化です。検索クエリには対応していません。したがって、RAGパイプラインを構築するための関連文書を検索するためには、Milvusベクターストアに問い合わせるために、MilvusClient インスタンス(milvus_client)を直接使用する必要があります。

Milvusから関連文書を検索する関数を定義する。

def retrieve_documents(question, top_k=3):
    query_vector = list(
        embed_anything.embed_query([question], embedder=model)[0].embedding
    )
    search_res = milvus_client.search(
        collection_name=index_name,
        data=[query_vector],
        limit=top_k,
        output_fields=["text"],
    )
    docs = [(res["entity"]["text"], res["distance"]) for res in search_res[0]]
    return docs

RAGパイプラインで取得した文書を使用してレスポンスを生成する関数を定義する。

def generate_rag_response(question):
    retrieved_docs = retrieve_documents(question)
    context = "\n".join([f"Text: {doc[0]}\n" for doc in retrieved_docs])
    system_prompt = (
        "You are an AI assistant. Provide answers based on the given context."
    )
    user_prompt = f"""
    Use the following pieces of information to answer the question. If the information is not in the context, say you don't know.
    
    Context:
    {context}
    
    Question: {question}
    """
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
    )
    return response.choices[0].message.content

サンプル質問でRAGパイプラインをテストしてみましょう。

question = "How does Milvus search for similar documents?"
answer = generate_rag_response(question)
print(f"Question: {question}")
print(f"Answer: {answer}")

Question: How does Milvus search for similar documents?
Answer: Milvus searches for similar documents primarily through Approximate Nearest Neighbor (ANN) search, which finds the top K vectors closest to a given query vector. It also supports various other types of searches, such as filtering search under specified conditions, range search within a specified radius, hybrid search based on multiple vector fields, and keyword search based on BM25. Additionally, it can perform reranking to adjust the order of search results based on additional criteria, refining the initial ANN search results.