GraphRagでグローバル検索システムの構築：AI駆動のデータ分析入門

はじめに
Docker Composeを使用した環境構築
環境設定
LLM（大規模言語モデル）のセットアップ
データの読み込み
グローバルコンテキストの構築
検索パラメータの設定
検索エンジンの構築
検索の実行
結果の分析
まとめ
リポジトリ
1. 関連

はじめに

このチュートリアルでは、大規模なデータセットに対して効率的な検索を行う「グローバル検索システム」の構築方法を学びます。AIが生成したコミュニティレポートを活用し、複雑な質問に対する回答を生成する方法を解説していきます。

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

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

Docker Composeを使用した環境構築

このプロジェクトではDocker Composeを使用して開発環境を構築しています。以下に、Dockerfileとdocker-compose.ymlファイルの内容を示します。

Dockerfile

FROM python:3.11-slim

WORKDIR /app

# 必要なシステムパッケージをインストール
RUN apt-get update && apt-get install -y \
    build-essential \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Poetryをインストール
RUN curl -sSL https://install.python-poetry.org | python3 -
ENV PATH="${PATH}:/root/.local/bin"

# Poetryの仮想環境を無効化（Dockerコンテナ内では不要）
RUN poetry config virtualenvs.create false

# Jupyter LabとJupytextをインストール
RUN pip install jupyterlab jupytext

# Jupyter Labを起動（アクセスキーなし）
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--allow-root", "--no-browser", "--NotebookApp.token=''", "--NotebookApp.password=''"]

このDockerfileでは、Python 3.11のslimイメージを基に、必要なシステムパッケージ、Poetry、Jupyter Lab、Jupytextをインストールしています。

docker-compose.yml

version: '3.8'

services:
  python-dev:
    build: .
    volumes:
      - .:/app
    ports:
      - "8888:8888"  # Jupyter Lab用
    environment:
      - JUPYTER_ENABLE_LAB=yes
    command: jupyter lab --ip=0.0.0.0 --allow-root --no-browser
    env_file:
      - .env

このdocker-compose.ymlファイルでは、Dockerfileをビルドし、ホストのカレントディレクトリをコンテナの/appディレクトリにマウントしています。また、Jupyter Labのポートを8888番でホストにマッピングし、環境変数ファイル(.env)を読み込んでいます。

環境の起動

プロジェクトディレクトリで以下のコマンドを実行することで、Docker Compose環境を起動できます：

docker-compose up --build

これにより、Jupyter Labが起動し、ブラウザからhttp://localhost:8888でアクセスできるようになります。

環境設定

まずは必要なライブラリをインポートし、環境を整えましょう。

%cd /app

!poetry install

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

import os
import sys

sys.path.append('..')

# グラフラグ（GraphRag）ライブラリからの特定モジュールのインポート
from graphrag.query.indexer_adapters import read_indexer_entities, read_indexer_reports
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.structured_search.global_search.community_context import GlobalCommunityContext
from graphrag.query.structured_search.global_search.search import GlobalSearch

# 環境変数からAPIキーとモデル名を取得
api_key = os.environ["OPENAI_API_KEY"]
llm_model = "gpt-4o"

このコードブロックでは、必要なライブラリとモジュールをインポートしています。osモジュールは環境変数の取得に、pandasはデータ操作に、tiktokenはトークン化に使用します。graphragライブラリからは、グローバル検索に必要な特定のクラスやメソッドをインポートしています。

LLM（大規模言語モデル）のセットアップ

次に、LLMのインスタンスを作成します。

# LLMインスタンスの作成
llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,  # OpenAIのAPIを使用。Azure OpenAIを使用する場合は OpenaiApiType.AzureOpenAI に変更
    max_retries=20,  # API呼び出しが失敗した場合の最大再試行回数
)

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

ここでは、OpenAIのAPIを使用するLLMインスタンスを作成しています。api_keyとmodelは環境変数から取得しています。max_retriesパラメータは、API呼び出しが失敗した場合の再試行回数を設定します。

また、tiktokenライブラリを使用してトークンエンコーダーを設定しています。これは、テキストをトークンに分割する際に使用されます。

データの読み込み

次に、必要なデータを読み込みます。

# データファイルのパスとテーブル名の設定
%cd ./examples_notebooks
INPUT_DIR = "./inputs/operation dulce"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"

# コミュニティレベルの設定（高いほど細かい粒度のコミュニティを使用）
COMMUNITY_LEVEL = 2

# Parquetファイルからデータを読み込む
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")

# インデクサーからエンティティとレポートを読み込む
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)

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

# レポートデータの先頭5行を表示
report_df.head()

このコードブロックでは、Parquetファイルからエンティティ、レポート、エンティティ埋め込みのデータを読み込んでいます。read_indexer_reportsとread_indexer_entities関数を使用して、これらのデータを処理し、後続の分析に使用できる形式に変換しています。

グローバルコンテキストの構築

データを読み込んだら、グローバルコンテキストを構築します。

# グローバルコミュニティコンテキストの構築
context_builder = GlobalCommunityContext(
    community_reports=reports,
    entities=entities,  # エンティティを使用しない場合はNoneに設定可能
    token_encoder=token_encoder,
)

GlobalCommunityContextクラスを使用して、コミュニティレポートとエンティティからグローバルコンテキストを構築しています。このコンテキストは、後続の検索プロセスで使用されます。

検索パラメータの設定

グローバル検索を実行する前に、いくつかのパラメータを設定する必要があります。

# コンテキストビルダーのパラメータ設定
context_builder_params = {
    "use_community_summary": False,  # 完全なコミュニティレポートを使用（True: コミュニティの短い要約を使用）
    "shuffle_data": True,  # データをシャッフルするかどうか
    "include_community_rank": True,  # コミュニティランクを含めるかどうか
    "min_community_rank": 0,  # 最小コミュニティランク
    "community_rank_name": "rank",  # コミュニティランクの列名
    "include_community_weight": True,  # コミュニティの重みを含めるかどうか
    "community_weight_name": "occurrence weight",  # コミュニティの重みの列名
    "normalize_community_weight": True,  # コミュニティの重みを正規化するかどうか
    "max_tokens": 12_000,  # 使用する最大トークン数（モデルの制限に応じて調整）
    "context_name": "Reports",  # コンテキストの名前
}

# マッピング段階のLLMパラメータ
map_llm_params = {
    "max_tokens": 1000,  # 生成する最大トークン数
    "temperature": 0.0,  # 生成の多様性（0.0: 決定論的、1.0: 最大の多様性）
    "response_format": {"type": "json_object"},  # レスポンスのフォーマット
}

# リデュース段階のLLMパラメータ
reduce_llm_params = {
    "max_tokens": 2000,  # 生成する最大トークン数（モデルの制限に応じて調整）
    "temperature": 0.0,  # 生成の多様性
}

これらのパラメータは、検索プロセスの様々な側面を制御します。例えば、use_community_summaryをFalseに設定することで、完全なコミュニティレポートを使用するようになります。max_tokensパラメータは、使用するLLMモデルのトークン制限に応じて調整する必要があります。

検索エンジンの構築

パラメータを設定したら、グローバル検索エンジンを構築します。

# グローバル検索エンジンの構築
search_engine = GlobalSearch(
    llm=llm,  # 使用するLLMインスタンス
    context_builder=context_builder,  # コンテキストビルダー
    token_encoder=token_encoder,  # トークンエンコーダー
    max_data_tokens=12_000,  # 使用する最大データトークン数（モデルの制限に応じて調整）
    map_llm_params=map_llm_params,  # マッピング段階のLLMパラメータ
    reduce_llm_params=reduce_llm_params,  # リデュース段階のLLMパラメータ
    allow_general_knowledge=False,  # 一般知識の使用を許可するかどうか
    json_mode=True,  # JSONモードを使用するかどうか
    context_builder_params=context_builder_params,  # コンテキストビルダーのパラメータ
    concurrent_coroutines=32,  # 並行して実行するコルーチンの数
    response_type="multiple paragraphs",  # レスポンスタイプ（例：複数段落、リスト、単一段落など）
)

GlobalSearchクラスを使用して検索エンジンを構築します。ここでは、先ほど設定したLLM、コンテキストビルダー、各種パラメータを渡しています。allow_general_knowledgeをFalseに設定することで、モデルは与えられたコンテキストのみを使用して回答を生成します。

検索の実行

検索エンジンの準備ができたら、実際に検索を実行します。

# 非同期で検索を実行
result = await search_engine.asearch(
    "この物語の主な対立は何で、主人公と敵対者は誰ですか？"
)

# 検索結果を表示
print(result.response)

asearchメソッドを使用して非同期で検索を実行します。質問文を引数として渡し、結果をresult変数に格納します。result.responseには、LLMが生成した回答が含まれています。

結果の分析

最後に、検索結果を分析します。

# LLMレスポンスのコンテキスト構築に使用されたデータを検査
print(result.context_data["reports"])

# LLM呼び出し回数とトークン数を表示
print(f"LLM呼び出し回数: {result.llm_calls}. LLMトークン数: {result.prompt_tokens}")

result.context_data["reports"]を表示することで、LLMがレスポンスを生成する際に使用したコンテキストデータを確認できます。また、result.llm_callsとresult.prompt_tokensを表示することで、検索プロセス中のLLM呼び出し回数と使用されたトークン数を確認できます。