LocalSearchを使った高度な文書検索と質問応答システムの構築

はじめに
必要なライブラリのインポート
設定と定数の定義
エンティティデータの読み込みと処理
リレーションシップデータの読み込み
共変量データの読み込み
コミュニティレポートの読み込み
テキストユニットの読み込み
LLMとエンベディングモデルの設定
ローカル検索コンテキストビルダーの作成
ローカル検索エンジンの作成
ローカル検索の実行
検索結果のコンテキストデータの確認
質問生成
まとめ
リポジトリ
1. 関連

はじめに

この記事では、Microsoftが開発したGraphRAGライブラリを使用して、高度な文書検索と質問応答システムを構築する方法を解説します。LocalSearchという手法を用いて、AI抽出された知識グラフと生のドキュメントのテキストチャンクを組み合わせて回答を生成します。

この手法は、ドキュメント内で言及されている特定のエンティティに関する理解が必要な質問（例：「カモミールの癒し効果は何ですか？」）に適しています。

GraphRAGの検索手法：LocalSearchとGlobal Searchの徹底比較

はじめにみなさん、こんにちは！今日は、GraphRAGというすごく便利なライブラリの中にある二つの検索方法について、わかりやすくお話しします。その二つとは、「LocalSearch（ローカルサーチ）」と「Global Search（グローバ...

必要なライブラリのインポート

まずは、必要なライブラリをインポートします。

# 必要なライブラリをインポート
import os
import pandas as pd
import tiktoken

# GraphRAGの各種モジュールをインポート
from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
    read_indexer_covariates,
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.input.loaders.dfs import store_entity_semantic_embeddings
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import LocalSearchMixedContext
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores.lancedb import LanceDBVectorStore

これらのライブラリとモジュールは、データの読み込み、エンベディングの生成、検索エンジンの構築などに使用されます。

設定と定数の定義

次に、入力ディレクトリやテーブル名などの定数を定義します。

# 入力ディレクトリとデータベースURIの設定
INPUT_DIR = "./inputs/operation dulce"
LANCEDB_URI = f"{INPUT_DIR}/lancedb"

# 各テーブル名の定義
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"

# コミュニティレベルの設定
COMMUNITY_LEVEL = 2

これらの定数は、後続の処理でデータの読み込みやデータベースへの接続に使用されます。

エンティティデータの読み込みと処理

エンティティ（ノード）データを読み込み、処理します。

# エンティティデータの読み込み
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")

# エンティティオブジェクトの作成
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)

# LanceDBベクトルストアの初期化と接続
description_embedding_store = LanceDBVectorStore(
    collection_name="entity_description_embeddings",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)

# エンティティの意味的エンベディングの保存
entity_description_embeddings = store_entity_semantic_embeddings(
    entities=entities, vectorstore=description_embedding_store
)

# エンティティ数の表示
print(f"エンティティ数: {len(entity_df)}")
entity_df.head()

このコードブロックでは、エンティティデータを読み込み、LanceDBベクトルストアに保存しています。これにより、後続の検索処理で高速なベクトル検索が可能になります。

リレーションシップデータの読み込み

エンティティ間の関係性を表すリレーションシップデータを読み込みます。

# リレーションシップデータの読み込み
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)

# リレーションシップ数の表示
print(f"リレーションシップ数: {len(relationship_df)}")
relationship_df.head()

リレーションシップデータは、エンティティ間のつながりを表現し、より深い文脈理解に役立ちます。

共変量データの読み込み

共変量（この場合、主張）データを読み込みます。

# 共変量データの読み込み
covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet")
claims = read_indexer_covariates(covariate_df)

# 主張レコード数の表示
print(f"主張レコード数: {len(claims)}")
covariates = {"claims": claims}

共変量データは、エンティティに関連する追加情報を提供し、より豊かな文脈を構築するのに役立ちます。

コミュニティレポートの読み込み

コミュニティレポートデータを読み込みます。

# コミュニティレポートの読み込み
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)

# レポートレコード数の表示
print(f"レポートレコード数: {len(report_df)}")
report_df.head()

コミュニティレポートは、関連するエンティティのグループに関する要約情報を提供します。

テキストユニットの読み込み

テキストユニット（文書の断片）を読み込みます。

# テキストユニットの読み込み
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)

# テキストユニットレコード数の表示
print(f"テキストユニットレコード数: {len(text_unit_df)}")
text_unit_df.head()

テキストユニットは、元の文書から抽出された意味のある断片で、詳細な情報を提供します。

LLMとエンベディングモデルの設定

言語モデル（LLM）とエンベディングモデルを設定します。

# 環境変数から API キーとモデル名を取得
api_key = os.environ["GRAPHRAG_API_KEY"]
llm_model = os.environ["GRAPHRAG_LLM_MODEL"]
embedding_model = os.environ["GRAPHRAG_EMBEDDING_MODEL"]

# ChatOpenAI インスタンスの作成
llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,  # OpenAI API または Azure OpenAI API を選択
    max_retries=20,
)

# トークンエンコーダーの取得
token_encoder = tiktoken.get_encoding("cl100k_base")

# テキストエンベッダーの作成
text_embedder = OpenAIEmbedding(
    api_key=api_key,
    api_base=None,
    api_type=OpenaiApiType.OpenAI,
    model=embedding_model,
    deployment_name=embedding_model,
    max_retries=20,
)

これらの設定により、テキストの処理や意味的な類似性の計算が可能になります。

ローカル検索コンテキストビルダーの作成

ローカル検索のためのコンテキストビルダーを作成します。

# ローカル検索コンテキストビルダーの作成
context_builder = LocalSearchMixedContext(
    community_reports=reports,
    text_units=text_units,
    entities=entities,
    relationships=relationships,
    covariates=covariates,
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,
    text_embedder=text_embedder,
    token_encoder=token_encoder,
)

このコンテキストビルダーは、様々なデータソースを組み合わせて、質問に対する適切なコンテキストを構築します。

ローカル検索エンジンの作成

ローカル検索エンジンを作成し、パラメータを設定します。

# ローカル検索エンジンのパラメータ設定
local_context_params = {
    "text_unit_prop": 0.5,
    "community_prop": 0.1,
    "conversation_history_max_turns": 5,
    "conversation_history_user_turns_only": True,
    "top_k_mapped_entities": 10,
    "top_k_relationships": 10,
    "include_entity_rank": True,
    "include_relationship_weight": True,
    "include_community_rank": False,
    "return_candidate_context": False,
    "embedding_vectorstore_key": EntityVectorStoreKey.ID,
    "max_tokens": 12_000,
}

llm_params = {
    "max_tokens": 2_000,
    "temperature": 0.0,
}

# ローカル検索エンジンの作成
search_engine = LocalSearch(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    llm_params=llm_params,
    context_builder_params=local_context_params,
    response_type="multiple paragraphs",
)

これらのパラメータにより、検索エンジンの動作を細かく制御できます。

ローカル検索の実行

作成した検索エンジンを使用して、サンプルクエリに対する検索を実行します。

# サンプルクエリに対する検索の実行
result = await search_engine.asearch("Agent Mercerについて教えてください")
print(result.response)

question = "Dr. Jordan Hayesについて教えてください"
result = await search_engine.asearch(question)
print(result.response)

この例では、"Agent Mercer"と"Dr. Jordan Hayes"に関する情報を検索しています。

検索結果のコンテキストデータの確認

検索結果に使用されたコンテキストデータを確認します。

# エンティティデータの確認
result.context_data["entities"].head()

# リレーションシップデータの確認
result.context_data["relationships"].head()

# レポートデータの確認
result.context_data["reports"].head()

# ソースデータの確認
result.context_data["sources"].head()

# 主張データの確認（存在する場合）
if "claims" in result.context_data:
    print(result.context_data["claims"].head())

これらのデータを確認することで、検索エンジンがどのような情報を使用して回答を生成したかを理解できます。

質問生成

ユーザーの質問履歴から次の候補質問を生成する機能を実装します。

# 質問生成器の作成
question_generator = LocalQuestionGen(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    llm_params=llm_params,
    context_builder_params=local_context_params,
)

# 質問履歴の定義
question_history = [
    "Agent Mercerについて教えてください",
    "Dulce軍事基地で何が起こっていますか？",
]

# 候補質問の生成
candidate_questions = await question_generator.agenerate(
    question_history=question_history, context_data=None, question_count=5
)
print(candidate_questions.response)

この機能により、ユーザーとのインタラクティブな対話を支援し、より深い探索を促すことができます。