Milvus ハイブリッド・サーチ・レトリバー

概要

Milvusはオープンソースのベクターデータベースであり、類似検索やAIアプリケーションの組み込みに威力を発揮します。Milvusは非構造化データ検索をより身近なものにし、導入環境に関わらず一貫したユーザーエクスペリエンスを提供します。

これは、密なベクトル検索と疎なベクトル検索の両方の長所を組み合わせたMilvusハイブリッド検索レトリバーを使い始めるのに役立ちます。MilvusCollectionHybridSearchRetriever の全機能と設定に関する詳細なドキュメントはAPIリファレンスをご覧ください。

Milvus Multi-Vector Searchのドキュメントもご参照ください。

統合の詳細

レトリバー	セルフホスト	クラウド	パッケージ
MilvusCollectionHybridSearchRetriever	✅	❌	langchain_milvus

セットアップ

個々のクエリから自動トレースを取得したい場合は、以下のコメントを外してLangSmithAPIキーを設定することもできます：

# os.environ["LANGSMITH_API_KEY"] = getpass.getpass("Enter your LangSmith API key: ")
# os.environ["LANGSMITH_TRACING"] = "true"

インストール

このリトリーバはlangchain-milvus パッケージに含まれています。このガイドでは以下の依存関係を必要とします：

%pip install --upgrade --quiet pymilvus[model] langchain-milvus langchain-openai

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_milvus.retrievers import MilvusCollectionHybridSearchRetriever
from langchain_milvus.utils.sparse import BM25SparseEmbedding
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    WeightedRanker,
    connections,
)

Milvusサービスの開始

milvusサービスを開始するには、milvusのドキュメントを参照してください。

milvusを起動した後、milvus接続URIを指定する必要があります。

CONNECTION_URI = "http://localhost:19530"

OpenAI API Keyの準備

OpenAIのドキュメントを参照してOpenAIのAPIキーを取得し、環境変数に設定してください。

export OPENAI_API_KEY=<your_api_key>

密埋め込み関数と疎埋め込み関数の準備

10個の小説の偽の記述をフィクションにしてみましょう。実際の制作では、大量のテキストデータになる可能性があります。

texts = [
    "In 'The Whispering Walls' by Ava Moreno, a young journalist named Sophia uncovers a decades-old conspiracy hidden within the crumbling walls of an ancient mansion, where the whispers of the past threaten to destroy her own sanity.",
    "In 'The Last Refuge' by Ethan Blackwood, a group of survivors must band together to escape a post-apocalyptic wasteland, where the last remnants of humanity cling to life in a desperate bid for survival.",
    "In 'The Memory Thief' by Lila Rose, a charismatic thief with the ability to steal and manipulate memories is hired by a mysterious client to pull off a daring heist, but soon finds themselves trapped in a web of deceit and betrayal.",
    "In 'The City of Echoes' by Julian Saint Clair, a brilliant detective must navigate a labyrinthine metropolis where time is currency, and the rich can live forever, but at a terrible cost to the poor.",
    "In 'The Starlight Serenade' by Ruby Flynn, a shy astronomer discovers a mysterious melody emanating from a distant star, which leads her on a journey to uncover the secrets of the universe and her own heart.",
    "In 'The Shadow Weaver' by Piper Redding, a young orphan discovers she has the ability to weave powerful illusions, but soon finds herself at the center of a deadly game of cat and mouse between rival factions vying for control of the mystical arts.",
    "In 'The Lost Expedition' by Caspian Grey, a team of explorers ventures into the heart of the Amazon rainforest in search of a lost city, but soon finds themselves hunted by a ruthless treasure hunter and the treacherous jungle itself.",
    "In 'The Clockwork Kingdom' by Augusta Wynter, a brilliant inventor discovers a hidden world of clockwork machines and ancient magic, where a rebellion is brewing against the tyrannical ruler of the land.",
    "In 'The Phantom Pilgrim' by Rowan Welles, a charismatic smuggler is hired by a mysterious organization to transport a valuable artifact across a war-torn continent, but soon finds themselves pursued by deadly assassins and rival factions.",
    "In 'The Dreamwalker's Journey' by Lyra Snow, a young dreamwalker discovers she has the ability to enter people's dreams, but soon finds herself trapped in a surreal world of nightmares and illusions, where the boundaries between reality and fantasy blur.",
]

密なベクトルの生成にはOpenAI Embeddingを、疎なベクトルの生成にはBM25アルゴリズムを使います。

密埋め込み関数を初期化し、次元を取得する

dense_embedding_func = OpenAIEmbeddings()
dense_dim = len(dense_embedding_func.embed_query(texts[1]))
dense_dim

1536

スパース埋め込み関数を初期化する。

なお、スパース埋め込み関数の出力は、入力テキストのキーワードのインデックスと重みを表すスパースベクトルの集合である。

sparse_embedding_func = BM25SparseEmbedding(corpus=texts)
sparse_embedding_func.embed_query(texts[1])

{0: 0.4270424944042204,
 21: 1.845826690498331,
 22: 1.845826690498331,
 23: 1.845826690498331,
 24: 1.845826690498331,
 25: 1.845826690498331,
 26: 1.845826690498331,
 27: 1.2237754316221157,
 28: 1.845826690498331,
 29: 1.845826690498331,
 30: 1.845826690498331,
 31: 1.845826690498331,
 32: 1.845826690498331,
 33: 1.845826690498331,
 34: 1.845826690498331,
 35: 1.845826690498331,
 36: 1.845826690498331,
 37: 1.845826690498331,
 38: 1.845826690498331,
 39: 1.845826690498331}

Milvusコレクションの作成とデータのロード

接続URIの初期化と接続の確立

connections.connect(uri=CONNECTION_URI)

フィールド名とデータ型を定義する

pk_field = "doc_id"
dense_field = "dense_vector"
sparse_field = "sparse_vector"
text_field = "text"
fields = [
    FieldSchema(
        name=pk_field,
        dtype=DataType.VARCHAR,
        is_primary=True,
        auto_id=True,
        max_length=100,
    ),
    FieldSchema(name=dense_field, dtype=DataType.FLOAT_VECTOR, dim=dense_dim),
    FieldSchema(name=sparse_field, dtype=DataType.SPARSE_FLOAT_VECTOR),
    FieldSchema(name=text_field, dtype=DataType.VARCHAR, max_length=65_535),
]

定義されたスキーマでコレクションを作成

schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
collection = Collection(
    name="IntroductionToTheNovels", schema=schema, consistency_level="Strong"
)

密なベクトルと疎なベクトルのインデックスを定義する

dense_index = {"index_type": "FLAT", "metric_type": "IP"}
collection.create_index("dense_vector", dense_index)
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
collection.create_index("sparse_vector", sparse_index)
collection.flush()

エンティティをコレクションに挿入し、コレクションをロードする

entities = []
for text in texts:
    entity = {
        dense_field: dense_embedding_func.embed_documents([text])[0],
        sparse_field: sparse_embedding_func.embed_documents([text])[0],
        text_field: text,
    }
    entities.append(entity)
collection.insert(entities)
collection.load()

インスタンス化

これで、疎フィールドと密フィールドの検索パラメータを定義して、Retriever をインスタンス化できます：

sparse_search_params = {"metric_type": "IP"}
dense_search_params = {"metric_type": "IP", "params": {}}
retriever = MilvusCollectionHybridSearchRetriever(
    collection=collection,
    rerank=WeightedRanker(0.5, 0.5),
    anns_fields=[dense_field, sparse_field],
    field_embeddings=[dense_embedding_func, sparse_embedding_func],
    field_search_params=[dense_search_params, sparse_search_params],
    top_k=3,
    text_field=text_field,
)

このRetrieverの入力パラメータでは、このCollectionの2つのフィールドでハイブリッド検索を実行するために、密な埋め込みと疎な埋め込みを使用し、再ランク付けにWeightedRankerを使用する。最終的に、3つのトップKドキュメントが返される。

使用法

retriever.invoke("What are the story about ventures?")

[Document(page_content="In 'The Lost Expedition' by Caspian Grey, a team of explorers ventures into the heart of the Amazon rainforest in search of a lost city, but soon finds themselves hunted by a ruthless treasure hunter and the treacherous jungle itself.", metadata={'doc_id': '449281835035545843'}),
 Document(page_content="In 'The Phantom Pilgrim' by Rowan Welles, a charismatic smuggler is hired by a mysterious organization to transport a valuable artifact across a war-torn continent, but soon finds themselves pursued by deadly assassins and rival factions.", metadata={'doc_id': '449281835035545845'}),
 Document(page_content="In 'The Dreamwalker's Journey' by Lyra Snow, a young dreamwalker discovers she has the ability to enter people's dreams, but soon finds herself trapped in a surreal world of nightmares and illusions, where the boundaries between reality and fantasy blur.", metadata={'doc_id': '449281835035545846'})]

チェーン内での使用

ChatOpenAIを初期化し、プロンプトテンプレートを定義します。

llm = ChatOpenAI()

PROMPT_TEMPLATE = """
Human: You are an AI assistant, and provides answers to questions by using fact based and statistical information when possible.
Use the following pieces of information to provide a concise answer to the question enclosed in <question> tags.

<context>
{context}
</context>

<question>
{question}
</question>

Assistant:"""

prompt = PromptTemplate(
    template=PROMPT_TEMPLATE, input_variables=["context", "question"]
)

ドキュメントをフォーマットする関数を定義する

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

retrieverと他のコンポーネントを使用してチェーンを定義します。

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

定義したチェーンを使用してクエリを実行します。

rag_chain.invoke("What novels has Lila written and what are their contents?")

"Lila Rose has written 'The Memory Thief,' which follows a charismatic thief with the ability to steal and manipulate memories as they navigate a daring heist and a web of deceit and betrayal."

コレクションを削除します。

collection.drop()

API リファレンス

すべてのMilvusCollectionHybridSearchRetriever 機能と設定についての詳細なドキュメントは、API リファレンスを参照ください。