Milvusハイブリッド検索レトリバー

ハイブリッド検索は異なる検索パラダイムの長所を組み合わせ、検索精度とロバスト性を高めます。密なベクトル検索と疎なベクトル検索、そして複数の密なベクトル検索ストラテジーの組み合わせの両方の機能を活用し、多様なクエリに対して包括的で正確な検索を保証します。

この図は、最も一般的なハイブリッド検索シナリオである、密＋疎ハイブリッド検索を示している。この場合、候補は意味ベクトル類似性と正確なキーワードマッチングの両方を用いて検索される。これらの方法からの結果はマージされ、再ランク付けされ、最終的な答えを生成するためにLLMに渡される。このアプローチは、精度と意味理解のバランスが取れており、多様なクエリーシナリオに非常に効果的である。

密＋疎ハイブリッド検索に加えて、ハイブリッド戦略は複数の密ベクトルモデルを組み合わせることもできる。例えば、ある密なベクトルモデルは意味的なニュアンスを捉えることに特化し、別の密なベクトルモデルは文脈的な埋め込みやドメイン固有の表現に焦点を当てるかもしれない。これらのモデルからの結果をマージし、それらを再ランク付けすることで、このタイプのハイブリッド検索は、よりニュアンスがあり、コンテキストを考慮した検索プロセスを保証する。

LangChain Milvus統合はハイブリッド検索を実装する柔軟な方法を提供し、任意の数のmilvusベクトルフィールドと、任意のカスタム密埋込みモデルまたは疎埋込みモデルをサポートします。

このチュートリアルでは、最も一般的な密＋疎のケースから始め、一般的なハイブリッド検索の使い方を紹介します。

前提条件

このノートブックを実行する前に、以下の依存関係がインストールされていることを確認してください：

$ pip install --upgrade --quiet  langchain langchain-core langchain-community langchain-text-splitters langchain-milvus langchain-openai bs4 pymilvus[model] #langchain-voyageai

OpenAIのモデルを使います。OpenAIの環境変数OPENAI_API_KEY 。

import os

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

Milvus サーバURI (オプションでTOKEN) を指定してください。Milvusサーバのインストールと起動方法については、こちらのガイドを参照してください。

URI = "http://localhost:19530"
# TOKEN = ...

テーマやジャンルごとに分類された架空のストーリーをまとめたサンプルドキュメントを用意する。

from langchain_core.documents import Document

docs = [
    Document(
        page_content="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.",
        metadata={"category": "Mystery"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Post-Apocalyptic"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Heist/Thriller"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Science Fiction"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Science Fiction/Romance"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Fantasy"},
    ),
    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={"category": "Adventure"},
    ),
    Document(
        page_content="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.",
        metadata={"category": "Steampunk/Fantasy"},
    ),
    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={"category": "Adventure/Thriller"},
    ),
    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={"category": "Fantasy"},
    ),
]

密埋め込み＋疎埋め込み

オプション1（推奨）：密埋め込み＋Milvus BM25組み込み機能

密埋め込み＋Milvus BM25組み込み関数を使って、ハイブリッド検索ベクトルストアのインスタンスを組み立てる。

from langchain_milvus import Milvus, BM25BuiltInFunction
from langchain_openai import OpenAIEmbeddings


vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=OpenAIEmbeddings(),
    builtin_function=BM25BuiltInFunction(),  # output_field_names="sparse"),
    vector_field=["dense", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

上記のコードでは、BM25BuiltInFunction のインスタンスを定義し、Milvus オブジェクトに渡しています。BM25BuiltInFunction は、Milvus の Functionの軽量ラッパークラスです。密＋疎ハイブリッド検索Milvusベクトルストアのインスタンスを初期化するためにOpenAIEmbeddings 。

BM25BuiltInFunction Milvusはコーパスやトレーニングをクライアントに渡す必要がなく、Milvusサーバ側で自動的に処理されるため、ユーザは語彙やコーパスを気にする必要がない。さらに、ユーザーはアナライザーをカスタマイズして、BM25にカスタムテキスト処理を実装することもできる。

BM25BuiltInFunction の詳細についてはFull-Text-Searchと LangChainとmilvusを使った全文検索をご参照ください。

オプション2：高密度でカスタマイズされたLangChainスパース埋め込みを使う

langchain_milvus.utils.sparse からBaseSparseEmbedding クラスを継承し、embed_query とembed_documents メソッドを実装することで、スパース埋め込み処理をカスタマイズすることができます。これにより、項頻度統計(BM25など)やニューラルネットワーク(SPADEなど)に基づくスパース埋め込み方法をカスタマイズすることができます。

以下に例を示します：

from typing import Dict, List
from langchain_milvus.utils.sparse import BaseSparseEmbedding


class MyCustomEmbedding(BaseSparseEmbedding):  # inherit from BaseSparseEmbedding
    def __init__(self, model_path): ...  # code to init or load model

    def embed_query(self, query: str) -> Dict[int, float]:
        ...  # code to embed query
        return {  # fake embedding result
            1: 0.1,
            2: 0.2,
            3: 0.3,
            # ...
        }

    def embed_documents(self, texts: List[str]) -> List[Dict[int, float]]:
        ...  # code to embed documents
        return [  # fake embedding results
            {
                1: 0.1,
                2: 0.2,
                3: 0.3,
                # ...
            }
        ] * len(texts)

langchain_milvus.utils.sparse にBaseSparseEmbedding から継承したデモクラスBM25SparseEmbedding があります。 他のlangchain密埋め込みクラスと同様に、Milvusベクトルストアのインスタンスの初期化埋め込みリストに渡すことができます。

# BM25SparseEmbedding is inherited from BaseSparseEmbedding
from langchain_milvus.utils.sparse import BM25SparseEmbedding

embedding1 = OpenAIEmbeddings()

corpus = [doc.page_content for doc in docs]
embedding2 = BM25SparseEmbedding(
    corpus=corpus
)  # pass in corpus to initialize the statistics

vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=[embedding1, embedding2],
    vector_field=["dense", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

これはBM25を使う方法ですが、用語頻度統計のためにコーパスを管理する必要があります。BM25はMilvusサーバ側で全ての処理を行うため、BM25組み込み関数(Option 1)を使用することをお勧めします。これにより、ユーザはコーパスの管理や語彙の学習に煩わされることがなくなります。詳しくはLangChainとMilvusによる全文検索をご参照ください。

複数の任意ベクトルフィールドの定義

Milvusのベクトルストアを初期化する際に、埋込み候補のリスト（将来的にはビルドイン関数のリストも）を渡すことで、複数の埋込み候補を検索し、再ランク付けすることができます。 以下に例を示します：

# from langchain_voyageai import VoyageAIEmbeddings

embedding1 = OpenAIEmbeddings(model="text-embedding-ada-002")
embedding2 = OpenAIEmbeddings(model="text-embedding-3-large")
# embedding3 = VoyageAIEmbeddings(model="voyage-3")  # You can also use embedding from other embedding model providers, e.g VoyageAIEmbeddings


vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=[embedding1, embedding2],  # embedding3],
    builtin_function=BM25BuiltInFunction(output_field_names="sparse"),
    # `sparse` is the output field name of BM25BuiltInFunction, and `dense1` and `dense2` are the output field names of embedding1 and embedding2
    vector_field=["dense1", "dense2", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

vectorstore.vector_fields

['dense1', 'dense2', 'sparse']

この例では、3つのベクトル場があります。その中で、sparse はBM25BuiltInFunction の出力フィールドとして使用され、他の2つ、dense1 とdense2 は、（順番に基づき）2つのOpenAIEmbeddings モデルの出力フィールドとして自動的に割り当てられる。

マルチ・ベクトル・フィールドのインデックス・パラメータの指定

デフォルトでは、各ベクトル・フィールドのインデックス・タイプは、埋め込みまたは組み込み関数のタイプによって自動的に決定されます。しかし、検索性能を最適化するために、各ベクトル・フィールドのインデックス・タイプを指定することもできます。

dense_index_param_1 = {
    "metric_type": "COSINE",
    "index_type": "HNSW",
}
dense_index_param_2 = {
    "metric_type": "IP",
    "index_type": "HNSW",
}
sparse_index_param = {
    "metric_type": "BM25",
    "index_type": "AUTOINDEX",
}

vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=[embedding1, embedding2],
    builtin_function=BM25BuiltInFunction(output_field_names="sparse"),
    index_params=[dense_index_param_1, dense_index_param_2, sparse_index_param],
    vector_field=["dense1", "dense2", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

vectorstore.vector_fields

['dense1', 'dense2', 'sparse']

候補の再ランク付け

検索の第一段階の後、より良い結果を得るために候補を再ランク付けする必要があります。要件に応じてWeightedRankerまたはRRFRankerを選択することができます。詳しくはRerankingを参照してください。

以下は重み付け再ランク付けの例です：

vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=OpenAIEmbeddings(),
    builtin_function=BM25BuiltInFunction(),
    vector_field=["dense", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

query = "What are the novels Lila has written and what are their contents?"

vectorstore.similarity_search(
    query, k=1, ranker_type="weighted", ranker_params={"weights": [0.6, 0.4]}
)

[Document(metadata={'pk': 454646931479252186, 'category': 'Heist/Thriller'}, page_content="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.")]

以下はRRFリランカーの例です：

vectorstore.similarity_search(query, k=1, ranker_type="rrf", ranker_params={"k": 100})

[Document(metadata={'category': 'Heist/Thriller', 'pk': 454646931479252186}, page_content="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.")]

再ランクに関するパラメータを渡さない場合、デフォルトで平均加重再ランク戦略が使用されます。

RAGにおけるハイブリッド検索と再ランクの使用

RAGのシナリオでは、ハイブリッド検索の最も一般的なアプローチは、密検索＋疎検索であり、その後に再ランク付けが続く。次の例は、エンド・ツー・エンドの簡単なコードを示している。

データの準備

Langchain WebBaseLoaderを使ってウェブソースからドキュメントをロードし、RecursiveCharacterTextSplitterを使ってチャンクに分割する。

import bs4
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Create a WebBaseLoader instance to load documents from web sources
loader = WebBaseLoader(
    web_paths=(
        "https://lilianweng.github.io/posts/2023-06-23-agent/",
        "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    ),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(
            class_=("post-content", "post-title", "post-header")
        )
    ),
)
# Load documents from web sources using the loader
documents = loader.load()
# Initialize a RecursiveCharacterTextSplitter for splitting text into chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)

# Split the documents into chunks using the text_splitter
docs = text_splitter.split_documents(documents)

# Let's take a look at the first document
docs[1]

Document(metadata={'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/'}, page_content='Fig. 1. Overview of a LLM-powered autonomous agent system.\nComponent One: Planning#\nA complicated task usually involves many steps. An agent needs to know what they are and plan ahead.\nTask Decomposition#\nChain of thought (CoT; Wei et al. 2022) has become a standard prompting technique for enhancing model performance on complex tasks. The model is instructed to “think step by step” to utilize more test-time computation to decompose hard tasks into smaller and simpler steps. CoT transforms big tasks into multiple manageable tasks and shed lights into an interpretation of the model’s thinking process.\nTree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each step. It first decomposes the problem into multiple thought steps and generates multiple thoughts per step, creating a tree structure. The search process can be BFS (breadth-first search) or DFS (depth-first search) with each state evaluated by a classifier (via a prompt) or majority vote.\nTask decomposition can be done (1) by LLM with simple prompting like "Steps for XYZ.\\n1.", "What are the subgoals for achieving XYZ?", (2) by using task-specific instructions; e.g. "Write a story outline." for writing a novel, or (3) with human inputs.\nAnother quite distinct approach, LLM+P (Liu et al. 2023), involves relying on an external classical planner to do long-horizon planning. This approach utilizes the Planning Domain Definition Language (PDDL) as an intermediate interface to describe the planning problem. In this process, LLM (1) translates the problem into “Problem PDDL”, then (2) requests a classical planner to generate a PDDL plan based on an existing “Domain PDDL”, and finally (3) translates the PDDL plan back into natural language. Essentially, the planning step is outsourced to an external tool, assuming the availability of domain-specific PDDL and a suitable planner which is common in certain robotic setups but not in many other domains.\nSelf-Reflection#')

Milvusベクトルストアにドキュメントをロードする。

Milvusベクターストアには2つのベクターフィールドがあります。dense はOpenAIのエンベッディング用、sparse はBM25関数用です。

vectorstore = Milvus.from_documents(
    documents=docs,
    embedding=OpenAIEmbeddings(),
    builtin_function=BM25BuiltInFunction(),
    vector_field=["dense", "sparse"],
    connection_args={
        "uri": URI,
    },
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/tune_consistency.md#Consistency-Level for more details.
    drop_old=False,
)

RAGチェーンの構築

LLMインスタンスとプロンプトを準備し、LangChain Expression Languageを使ってRAGパイプラインに結合する。

from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

# Initialize the OpenAI language model for response generation
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)

# Define the prompt template for generating AI responses
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.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
<context>
{context}
</context>

<question>
{question}
</question>

The response should be specific and use statistics or numbers when possible.

Assistant:"""

# Create a PromptTemplate instance with the defined template and input variables
prompt = PromptTemplate(
    template=PROMPT_TEMPLATE, input_variables=["context", "question"]
)
# Convert the vector store to a retriever
retriever = vectorstore.as_retriever()


# Define a function to format the retrieved documents
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

LCEL(LangChain Expression Language)を使ってRAGチェインを構築する。

# Define the RAG (Retrieval-Augmented Generation) chain for AI response generation
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# rag_chain.get_graph().print_ascii()

特定の質問でRAGチェーンを呼び出し、応答を取得する。

query = "What is PAL and PoT?"
res = rag_chain.invoke(query)
res

'PAL (Program-aided Language models) and PoT (Program of Thoughts prompting) are approaches that involve using language models to generate programming language statements to solve natural language reasoning problems. This method offloads the solution step to a runtime, such as a Python interpreter, allowing for complex computation and reasoning to be handled externally. PAL and PoT rely on language models with strong coding skills to effectively perform these tasks.'

おめでとう！MilvusとLangChainのハイブリッド（密なベクトル＋疎なbm25関数）検索RAGチェーンが構築できました。