Open8

LlamaIndexのGraphRAG実装を試してみる

kun432kun432

GraphRAGそのもののコンセプトについては以下を参考にした。

概念というかイメージを掴むのにわかりやすかった。
https://zenn.dev/zawawahoge/scraps/5a376b1c419cda

自分も、コスト感とか実現するまでの道のりがちょっとしんどくないかなという所感はあった。
https://note.com/avakansai/n/n82dcf5c8d02c

20240729追記

全体像掴むだけならこれが一番わかりやすいかも

https://internet.watch.impress.co.jp/docs/column/shimizu/1608736.html

kun432kun432

LlamaIndexでやる前に、まずGraphRAG公式のGetting Startedを試してみた。Colaboratoryで。

https://microsoft.github.io/graphrag/posts/get_started/

パッケージインストール

!pip install graphrag

サンプルドキュメントを用意。以下を使用。

https://ja.wikipedia.org/wiki/オグリキャップ

!mkdir -p ./ragtest/input
from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["オグリキャップ"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n## 概要\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("ragtest/input")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.txt", "w") as fp:
        fp.write(wiki_text)

設定ファイル等を生成

!python -m graphrag.index --init --root ./ragtest

こんな感じで生成された。

$ tree ragtest
ragtest
├── input
│   └── オグリキャップ.txt
├── output
│   └── 20240731-034904
│       └── reports
│           └── indexing-engine.log
├── prompts
│   ├── claim_extraction.txt
│   ├── community_report.txt
│   ├── entity_extraction.txt
│   └── summarize_descriptions.txt
└── settings.yaml

5 directories, 7 files

settings.yamlは以下

settings.yaml
encoding_model: cl100k_base
skip_workflows: []
llm:
  api_key: ${GRAPHRAG_API_KEY}
  type: openai_chat # or azure_openai_chat
  model: gpt-4-turbo-preview
  model_supports_json: true # recommended if this is available for your model.
  # max_tokens: 4000
  # request_timeout: 180.0
  # api_base: https://<instance>.openai.azure.com
  # api_version: 2024-02-15-preview
  # organization: <organization_id>
  # deployment_name: <azure_model_deployment_name>
  # tokens_per_minute: 150_000 # set a leaky bucket throttle
  # requests_per_minute: 10_000 # set a leaky bucket throttle
  # max_retries: 10
  # max_retry_wait: 10.0
  # sleep_on_rate_limit_recommendation: true # whether to sleep when azure suggests wait-times
  # concurrent_requests: 25 # the number of parallel inflight requests that may be made
  # temperature: 0 # temperature for sampling
  # top_p: 1 # top-p sampling
  # n: 1 # Number of completions to generate

parallelization:
  stagger: 0.3
  # num_threads: 50 # the number of threads to use for parallel processing

async_mode: threaded # or asyncio

embeddings:
  ## parallelization: override the global parallelization settings for embeddings
  async_mode: threaded # or asyncio
  llm:
    api_key: ${GRAPHRAG_API_KEY}
    type: openai_embedding # or azure_openai_embedding
    model: text-embedding-3-small
    # api_base: https://<instance>.openai.azure.com
    # api_version: 2024-02-15-preview
    # organization: <organization_id>
    # deployment_name: <azure_model_deployment_name>
    # tokens_per_minute: 150_000 # set a leaky bucket throttle
    # requests_per_minute: 10_000 # set a leaky bucket throttle
    # max_retries: 10
    # max_retry_wait: 10.0
    # sleep_on_rate_limit_recommendation: true # whether to sleep when azure suggests wait-times
    # concurrent_requests: 25 # the number of parallel inflight requests that may be made
    # batch_size: 16 # the number of documents to send in a single request
    # batch_max_tokens: 8191 # the maximum number of tokens to send in a single request
    # target: required # or optional
  


chunks:
  size: 1200
  overlap: 100
  group_by_columns: [id] # by default, we don't allow chunks to cross documents
    
input:
  type: file # or blob
  file_type: text # or csv
  base_dir: "input"
  file_encoding: utf-8
  file_pattern: ".*\\.txt$"

cache:
  type: file # or blob
  base_dir: "cache"
  # connection_string: <azure_blob_storage_connection_string>
  # container_name: <azure_blob_storage_container_name>

storage:
  type: file # or blob
  base_dir: "output/${timestamp}/artifacts"
  # connection_string: <azure_blob_storage_connection_string>
  # container_name: <azure_blob_storage_container_name>

reporting:
  type: file # or console, blob
  base_dir: "output/${timestamp}/reports"
  # connection_string: <azure_blob_storage_connection_string>
  # container_name: <azure_blob_storage_container_name>

entity_extraction:
  ## llm: override the global llm settings for this task
  ## parallelization: override the global parallelization settings for this task
  ## async_mode: override the global async_mode settings for this task
  prompt: "prompts/entity_extraction.txt"
  entity_types: [organization,person,geo,event]
  max_gleanings: 1

summarize_descriptions:
  ## llm: override the global llm settings for this task
  ## parallelization: override the global parallelization settings for this task
  ## async_mode: override the global async_mode settings for this task
  prompt: "prompts/summarize_descriptions.txt"
  max_length: 500

claim_extraction:
  ## llm: override the global llm settings for this task
  ## parallelization: override the global parallelization settings for this task
  ## async_mode: override the global async_mode settings for this task
  # enabled: true
  prompt: "prompts/claim_extraction.txt"
  description: "Any claims or facts that could be relevant to information discovery."
  max_gleanings: 1

community_reports:
  ## llm: override the global llm settings for this task
  ## parallelization: override the global parallelization settings for this task
  ## async_mode: override the global async_mode settings for this task
  prompt: "prompts/community_report.txt"
  max_length: 2000
  max_input_length: 8000

cluster_graph:
  max_cluster_size: 10

embed_graph:
  enabled: false # if true, will generate node2vec embeddings for nodes
  # num_walks: 10
  # walk_length: 40
  # window_size: 2
  # iterations: 3
  # random_seed: 597832

umap:
  enabled: false # if true, will generate UMAP embeddings for nodes

snapshots:
  graphml: false
  raw_entities: false
  top_level_nodes: false

local_search:
  # text_unit_prop: 0.5
  # community_prop: 0.1
  # conversation_history_max_turns: 5
  # top_k_mapped_entities: 10
  # top_k_relationships: 10
  # llm_temperature: 0 # temperature for sampling
  # llm_top_p: 1 # top-p sampling
  # llm_n: 1 # Number of completions to generate
  # max_tokens: 12000

global_search:
  # llm_temperature: 0 # temperature for sampling
  # llm_top_p: 1 # top-p sampling
  # llm_n: 1 # Number of completions to generate
  # max_tokens: 12000
  # data_max_tokens: 12000
  # map_max_tokens: 1000
  # reduce_max_tokens: 2000
  # concurrency: 32

モデルだけ書き換えた。LLMはgpt-4o、Embeddingはtext-embedding-3-largeで。

settings.yaml
(snip)
llm:
  api_key: ${GRAPHRAG_API_KEY}
  type: openai_chat # or azure_openai_chat
  model: gpt-4o
  model_supports_json: true
(snip)
settings.yaml
(snip)
embeddings:
  ## parallelization: override the global parallelization settings for embeddings
  async_mode: threaded # or asyncio
  llm:
    api_key: ${GRAPHRAG_API_KEY}
    type: openai_embedding # or azure_openai_embedding
    model: text-embedding-3-large
(snip)

上にもあるが、OpenAIのAPIキーを環境変数GRAPHRAG_API_KEYとしてセットする

import os
from google.colab import userdata

os.environ["GRAPHRAG_API_KEY"] = userdata.get('OPENAI_API_KEY')

インデックスを作成

!python -m graphrag.index --root ./ragtest

ではまずグローバルサーチ。

!python -m graphrag.query \
    --root ./ragtest \
    --method global \
    "何について書かれている文章ですか?"
INFO: Reading settings from ragtest/settings.yaml
creating llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_chat", 'model': 'gpt-4o', 'max_tokens': 4000, 'temperature': 0.0, 'top_p': 1.0, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': True, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}

SUCCESS: Global Search Response: ## 文章の内容について

文章は主に日本の競馬界と有名な競走馬オグリキャップに関する内容です。以下のセクションで、オグリキャップの影響力、競走成績、関係者、そしてその遺産について詳述されています。

### オグリキャップの影響力と競走成績

オグリキャップは日本の競馬界において非常に影響力のある競走馬であり、その競走成績は特に有馬記念や天皇賞(秋)などの主要なレースで際立っています [Data: Reports (31, 25, 22, 29, +more)]。彼の成功は中央競馬(JRA)にも大きな影響を与えました [Data: Reports (7, 32)]。

### 関係者と血統

オグリキャップに関連する重要な人物として、騎手やトレーナー、所有者などが挙げられます。特に小椋孝一や鷲見昌勇など、彼の血統や育成に関わった人物についても詳細に記述されています [Data: Reports (21, 26, 15, 8, +more)]。

### 人気とメディアの注目

オグリキャップの人気は非常に高く、メディアの注目も集めました。特に引退式や記念レースなどのイベントにおけるファンの反応が記載されています [Data: Reports (31, 16, 22)]。

### 遺産と子孫

オグリキャップの遺産はその子孫にも引き継がれており、例えばレディアイコやミンナノヒーローなどの競走成績についても触れられています [Data: Reports (1)]。

このように、文章はオグリキャップの多岐にわたる影響とその周囲のコミュニティについて詳細に述べています。

次にローカルサーチ

!python -m graphrag.query \
    --root ./ragtest \
    --method local \
    "オグリキャップの主な勝ち鞍は?"
INFO: Reading settings from ragtest/settings.yaml
[2024-07-31T04:00:06Z WARN  lance::dataset] No existing dataset at /content/lancedb/description_embedding.lance, it will be created
creating llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_chat", 'model': 'gpt-4o', 'max_tokens': 4000, 'temperature': 0.0, 'top_p': 1.0, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': True, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}
creating embedding llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_embedding", 'model': 'text-embedding-3-large', 'max_tokens': 4000, 'temperature': 0, 'top_p': 1, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': None, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}

SUCCESS: Local Search Response: # オグリキャップの主な勝ち鞍

オグリキャップはその卓越した競走能力で多くの主要レースを制し、日本の競馬史に名を刻みました。以下に、オグリキャップの主な勝ち鞍を紹介します。

## アリマ記念

オグリキャップはアリマ記念での勝利が特に有名です。1990年のアリマ記念では、体調不良にもかかわらず見事な走りを見せ、ファンの期待に応えました。この勝利はオグリキャップの初のGIレース勝利であり、その後の競馬ブームの象徴となりました [Data: Relationships (40, 79)]。

## 天皇賞(秋)

天皇賞(秋)もオグリキャップの重要な勝ち鞍の一つです。このレースでは、オグリキャップの強さと精神力が際立ちました。特に1988年の天皇賞(秋)では、オグリキャップは2位に終わりましたが、その後のレースでの活躍が期待されました [Data: Relationships (31)]。

## 毎日王冠

オグリキャップは毎日王冠でも勝利を収めています。このレースでの勝利により、オグリキャップはJRAの連続勝利記録に並び、その実力を証明しました [Data: Relationships (30)]。

## ニュージーランドトロフィー4歳ステークス

オグリキャップはニュージーランドトロフィー4歳ステークスでも圧勝し、その人気を一気に高めました。このレースでの勝利は、オグリキャップが競馬人気の旗手となるきっかけとなりました [Data: Entities (457)]。

## その他の主要レース

オグリキャップはその他にも多くの主要レースで勝利を収めています。例えば、Mainichi HaiやYasuda Kinenなどのレースでの勝利もその一部です [Data: Entities (0)]。

## 結論

オグリキャップの勝ち鞍はその競走能力と精神力を象徴するものであり、日本の競馬史において重要な位置を占めています。これらの勝利は、オグリキャップがいかに優れた競走馬であったかを物語っています。

いろいろハルシネーションしてるのと、ちょいちょいカタカナや英語になってるものがある。このあたりはプロンプトでそうなってしまうのだと思う。

この質問についてはグローバルでも聞いてみる。

!python -m graphrag.query \
    --root ./ragtest \
    --method global \
    "オグリキャップの主な勝ち鞍は?"
INFO: Reading settings from ragtest/settings.yaml
creating llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_chat", 'model': 'gpt-4o', 'max_tokens': 4000, 'temperature': 0.0, 'top_p': 1.0, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': True, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}

SUCCESS: Global Search Response: ### オグリキャップの主な勝ち鞍

オグリキャップは日本競馬界で非常に成功した競走馬であり、多くの主要なレースで勝利を収めました。以下に、彼の主な勝ち鞍をまとめます。

#### 有馬記念
オグリキャップは有馬記念で勝利を収めました。このレースは彼の人気と実力を示す重要な勝利であり、特に第35回有馬記念では、体調が悪い中での勝利が注目されました [Data: Reports (22, 31)]。

#### 中央競馬移籍後の主要勝利
1988年1月に中央競馬(JRA)に移籍してから、オグリキャップは12の主要な勝利を記録しました。これには、サツキショウ、東京優駿、安田記念などの著名なレースが含まれます [Data: Reports (7)]。

#### 天皇賞(秋)
オグリキャップは天皇賞(秋)でも勝利を収めました。このレースは彼の競争力と実力を示す重要な勝利の一つです [Data: Reports (23, 25)]。

#### 高松宮杯
高松宮杯でも勝利を収めたオグリキャップは、中央競馬に転入後初めて古馬と対戦し、コースレコードを樹立しました [Data: Reports (10)]。

#### 毎日王冠
オグリキャップは毎日王冠でも勝利を収め、連続重賞勝利のJRA記録に並びました。このレースは彼の卓越した能力を示すものでした [Data: Reports (29)]。

### その他のレース
オグリキャップはジャパンカップにも出走しましたが、1990年のジャパンカップでは11位に終わりました [Data: Reports (25)]。

オグリキャップのこれらの勝利は、彼の競走馬としての卓越した能力と競争力を示すものであり、日本競馬史において重要な位置を占めています。

どういう風にチャンク分割してグラフが生成されているかはわからないけども、元々の文章だと勝利については多分いくつかのパラグラフに分かれていたと思うので、こういう質問ならばグローバルサーチの場合でもあまり違いはなさそうかな。

別の質問をローカルサーチで。

!python -m graphrag.query \
    --root ./ragtest \
    --method local \
    "オグリキャップの血統は?"
INFO: Reading settings from ragtest/settings.yaml
creating llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_chat", 'model': 'gpt-4o', 'max_tokens': 4000, 'temperature': 0.0, 'top_p': 1.0, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': True, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}
creating embedding llm client with {'api_key': 'REDACTED,len=51', 'type': "openai_embedding", 'model': 'text-embedding-3-large', 'max_tokens': 4000, 'temperature': 0, 'top_p': 1, 'n': 1, 'request_timeout': 180.0, 'api_base': None, 'api_version': None, 'organization': None, 'proxy': None, 'cognitive_services_endpoint': None, 'deployment_name': None, 'model_supports_json': None, 'tokens_per_minute': 0, 'requests_per_minute': 0, 'max_retries': 10, 'max_retry_wait': 10.0, 'sleep_on_rate_limit_recommendation': True, 'concurrent_requests': 25}

SUCCESS: Local Search Response: # オグリキャップの血統

オグリキャップの血統は、競走馬としての成功に大きく寄与した要素の一つです。以下に、オグリキャップの血統に関する詳細を説明します。

## 父系と母系

### 父系
オグリキャップの父はダンシングキャップです。ダンシングキャップは短距離ダートレースに適した血統と評価されていましたが、オグリキャップのような大物を出すことは稀でした [Data: Entities (12); Relationships (139)]. ダンシングキャップの父系はネイティヴダンサーに遡り、ネイティヴダンサーは時折大物を出すことで知られています [Data: Entities (476); Relationships (169)]。

### 母系
オグリキャップの母はホワイトナルビーで、彼女は現役時代に笠松競馬場で4勝を挙げました。ホワイトナルビーの産駒は全て競馬で勝利を収めており、特に地方競馬で活躍しました [Data: Sources (20)]. さらに、ホワイトナルビーの5代母であるクインナルビーは1953年の天皇賞(秋)を制した名馬です [Data: Entities (477)]。

## 兄弟と子孫

オグリキャップにはいくつかの兄弟がいます。例えば、オグリローマンは1994年の桜花賞を優勝し、オグリイチバンとオグリトウショウは種牡馬として活躍しました [Data: Entities (481, 482); Relationships (172, 173)]。

オグリキャップの子孫も競馬界で一定の影響力を持っています。ノーザンキャップは中央競馬で3勝を挙げ、種牡馬としても活動しましたが、産駒は少数にとどまりました [Data: Entities (314); Relationships (119)]. クレイドルサイアーはノーザンキャップの唯一の産駒で、2013年に種牡馬登録されました [Data: Entities (318); Relationships (255)]。

## 血統の評価

オグリキャップの血統は、父系のダンシングキャップが「一発ある血統」と評される一方で、母系のホワイトナルビーの安定した成績も評価されています。血統評論家の山野浩一は、オグリキャップがネイティヴダンサーの隔世遺伝で生まれた可能性を指摘しています [Data: Entities (360)]。

オグリキャップの血統は、その競走成績とともに、競馬界における重要な遺産となっています。

こちらは良い感じ。

kun432kun432

Usage見てるとここまででだいたい$3ちょっとぐらいかな、ちょい高いと感じる。gpt-4o-miniにすれば安くはなるけど、試してみたら精度的に上よりも悪い感じだった、まあ当然かも。

kun432kun432

とりあえず

  • 仕組みをもう少しブレイクダウンして理解したい
  • プロンプトを日本語にしたい
  • グラフを可視化したい

プロンプトは英語のほうが精度高いのかもしれないけど