GraphRAG × Ollama ローカルでgraphRAG試してみた！

# はじめに
生成AIを業務に導入する際、RAGという手法がよく使われます。
RAGは企業の固有情報を元にカスタマイズされた回答を生成することが可能です。
今回はより精度の高い回答を得られると噂の「GraphRAG」をローカルで試してみました！

# GraphRAGとは
GraphRAGは、**知識グラフ**を使用してRAGの性能を向上させるアプローチです。
通常のRAGは**ベクトル検索**を用いて関連文書を取得し、回答を生成するのに対し、GraphRAGでは知識グラフを構築してデータ間の**関係を明確化**し、より高度な推論で回答を生成できます！

https://github.com/microsoft/graphrag

そんなGraphRAGは7月にOSSとして公開されています！

# Ollamaとは
Ollamaは、Llama 2などの大規模言語モデルをローカルで簡単に実行・管理できるオープンソースツールです。オープンソースLLMと組み合わせることでプライバシーを保護しつつ生成AIを利用することができます。

今回ローカルで試すにあたって「gemma2」と「nomic-embed-text」というオープンソースLLMを使っていきます！
こちらの二つはOllamaで用意されているので以下のコマンドを実行してモデルを用意しておきます。
``` bash
ollama pull gemma2
ollama pull nomic-embed-text
```

※Ollamaで用意されていないLLMを利用する場合のセットアップ方法については[以前書いた記事](https://qiita.com/yuta3okkie/items/29958d440dfb5d2abcfb)を参考にしてください！

# セットアップ
[公式ドキュメント](https://microsoft.github.io/graphrag/posts/get_started/)に従ってセットアップしていきます。

### 1\. graphRAGとollamaをインストール
以下のコマンドを実行しgraphragをインストールします。
``` bash
pip install graphrag
pip install ollama
```

### 2\. 初期セットアップコマンドを実行
任意のディレクトリに移動して以下のコマンドを実行します。
``` bash
python -m graphrag.index --init --root ./
```

このコマンドを実行することでカレントディレクトリ配下にgraphragを実行するのに必要なディレクトリとファイルが作成されます。

### 3\. inputディレクトリにデータソースとなるテキストファイルを配置
inputディレクトリを作成し、その中にテキストファイルを作成します。
今回はbook.txtというファイルに[桃太郎の英語版wiki](https://en.wikipedia.org/wiki/Momotar%C5%8D)のStoryを入れてデータソースとして利用しました。

### 3\. settings.yamlを更新
手順2で作成されたsettings.yamlを開き以下の部分を今回使用するモデルに合わせて修正します。
``` yaml
llm:
  api_key: ${GRAPHRAG_API_KEY}
  type: openai_chat # or azure_openai_chat
  model: gemma2
  model_supports_json: true # recommended if this is available for your model.
  # max_tokens: 4000
  # request_timeout: 180.0
  api_base: http://localhost:11434/v1
  # api_version: 2024-02-15-preview
```
``` yaml
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: nomic-embed-text
    api_base: http://localhost:11434/api
    # api_version: 2024-02-15-preview
```

### 4\. graphragのソースコードを一部カスタマイズ
このまま実行するとエラーになってしまうので以下二つのファイルをそれぞれ更新して保存する（[こちら](https://github.com/microsoft/graphrag/issues/345#issuecomment-2211819596)のissueを参考にさせていただきました！）
{Python格納ディレクトリ}/Lib/site-packages/graphrag/llm/openai/openai_embeddings_llm.py 30行目辺り _execute_llm内
``` python
async def _execute_llm(
    self, input: EmbeddingInput, **kwargs: Unpack[LLMInput]
) -> EmbeddingOutput | None:
    # args = {
    #     "model": self.configuration.model,
    #     **(kwargs.get("model_parameters") or {}),
    # }
    # embedding = await self.client.embeddings.create(
    #     input=input,
    #     **args,
    # )
    # return [d.embedding for d in embedding.data]
    embedding_list = []
    for inp in input:
        embedding = ollama.embeddings(model="nomic-embed-text", prompt=inp)
        embedding_list.append(embedding["embedding"])
    return embedding_list
```
{Python格納ディレクトリ}/Lib/site-packages/graphrag/llm/openai/utils.py 108行目辺り try_parse_json_object内
``` python
# _pattern = r"\{(.*)\}"
# _match = re.search(_pattern, input)
_pattern = r"\{(.*?)\}"
_match = re.search(_pattern, input, re.DOTALL)
input = "{" + _match.group(1) + "}" if _match else input
```

### 5\. graphragのindexerコマンドを実行
以下のコマンドを実行し、inputディレクトリに入ったデータソースから知識グラフを作成します。
``` bash 
python -m graphrag.index --root ./

 <途中省略>

⠴ GraphRAG Indexer
├── Loading Input (InputFileType.text) - 2 files loaded (0 filtered) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 0:00:00
├── create_base_text_units
├── create_base_extracted_entities
├── create_summarized_entities
├── create_base_entity_graph
├── create_final_entities
├── create_final_nodes
├── create_final_communities
├── join_text_units_to_entity_ids
├── create_final_relationships
├── join_text_units_to_relationship_ids
├── create_final_community_reports
├── create_final_text_units
├── create_base_documents
└── create_final_documents
🚀 All workflows completed successfully.
```

# 試してみる！

以下のコマンドを実行してこのストーリーの要約を聞いてみます。
``` bash
python -m graphrag.query --method global "What are the top themes in this story?"
```
出力結果がこちら
``` bash
INFO: Reading settings from settings.yaml
creating llm client with {'api_key': 'REDACTED,len=9', 'type': "openai_chat", 'model': 'gemma2', 'max_tokens': 4000, 'temperature': 0.0, 'top_p': 1.0, 'n': 1, 'request_timeout': 180.0, 'api_base': 'http://localhost:11434/v1', '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: The dataset appears to center around a narrative with clear thematic elements.

One prominent theme is the classic struggle between good and evil, exemplified by characters like Momotaro representing goodness and the Demons embodying wickedness. [Data: Reports (0)]


Let me know if you have any other questions about this dataset!
```
（和訳：顕著なテーマのひとつは、善と悪の間の古典的な闘争であり、善を象徴する桃太郎や悪を体現する鬼のようなキャラクターがその例である。）

それっぽい回答を返してくれました！！

# まとめ
ここまで読んでいただきありがとうございます！
今回はGraphRAGをOllamaを使ってローカルで試してみました！
オープンソース化されたことで一部カスタマイズをしてローカルでも試すことができました！！


最近は通常のベクトルRAGと今回のGraphRAGを合わせて使うHybridRAGというものが使われているそうです。
今後も引き続き生成AI周りの技術のキャッチアップを頑張ろうと思います！