はじめに
この記事では、Microsoftが開発したGraphRAGライブラリを使用して、高度な文書検索と質問応答システムを構築する方法を解説します。LocalSearchという手法を用いて、AI抽出された知識グラフと生のドキュメントのテキストチャンクを組み合わせて回答を生成します。
この手法は、ドキュメント内で言及されている特定のエンティティに関する理解が必要な質問（例：「カモミールの癒し効果は何ですか？」）に適しています。
必要なライブラリのインポート
まずは、必要なライブラリをインポートします。
# 必要なライブラリをインポート
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)この機能により、ユーザーとのインタラクティブな対話を支援し、より深い探索を促すことができます。
まとめ
この記事では、GraphRAGライブラリを使用して、高度な文書検索と質問応答システムを構築する方法を解説しました。LocalSearchメソッドを使用することで、AI抽出された知識グラフと生のドキュメントのテキストチャンクを組み合わせて、より的確な回答を生成できます。
このシステムは、特定のエンティティに関する深い理解が必要な質問に特に適しており、様々な分野での応用が期待できます。例えば、企業の内部文書検索、学術研究支援、カスタマーサポートなどに活用できるでしょう。
今後の展開としては、モデルの精度向上、多言語対応、リアルタイムデータ更新機能の追加などが考えられます。また、ユーザーインターフェースの改善や、より高度な対話機能の実装も興味深い課題となるでしょう。
この記事を参考に、皆さんも独自の高度な検索・質問応答システムの構築にチャレンジしてみてください。
コメント