話題のGraphRAGにAWSで挑戦しよう！（LlamaIndexとNeptuneに入門）

GraphRAGとは？

Microsoft社が、以前から提唱していたRAGの新しいアプローチ「GraphRAG」のサンプル実装を7月にGitHubで公開したことから、LLM界隈で一気に話題となりました。

従来のRAGは、検索対象のドキュメントを埋め込み（ベクトル）に変換しておくことで、検索クエリーと意味的に近いかたまり（チャンク）を数学的な計算によって探し当てて持ってくる手法でした。

これに対してGraphRAGでは、検索対象のドキュメントを埋め込みではなく「グラフ」として保存する手法です。これによって、従来のRAGでは検索精度がイマイチだったユースケースでも精度が改善することがあります。

グラフとは？

グラフといえば、Excelなどでデータを可視化するアレを思い浮かべる方も多いでしょう。GraphRAGで扱うのはそれではなく、数学の世界でいう「グラフ理論」のグラフです。ノード（節点）とエッジ（辺）で様々なものの関係を表す概念です。

これに対して、Excelなどで出てくる「棒グラフ」や「円グラフ」などは別物で、グラフではなく「チャート」と表現することもあります。

グラフには色々なものがありますが、GraphRAGで使うのは「知識グラフ（Knowledge graph）」です。社内データなどのナレッジをグラフとして保存するというわけですね。

知識グラフはどこに保存すればいい？

グラフを保存するためには、専用のデータベースである「グラフDB」が使われます。有名な製品に、OSSであるNeo4jなどがあります。（ちょうど先週、このNeo4jの社員らによる書籍の日本語版が出版されています）

AWSクラウドを利用する場合、マネージドなグラフDBサービスとしてAmazon Neptuneが利用できます。

Neptune Database

Amazon Neptuneに元からある機能が「データベース」です。Amazon RDSなどと同じように、クラスターを作成して利用します。従来のプロビジョンド版に加えて、2022年にサーバーレス版も追加されています。

※AWSあるあるで、サーバーレス版といいつつ最低稼働コストが発生するのですが、何ならプロビジョンド版を最低スペックで利用する方が安く済みます。（苦笑）

Neptune Analytics

2023年末のAWS re:Inventで発表された新機能です。大規模ワークロード向けとなっており、Neptune Databaseよりも膨大な（数百億もの接続をもつ）グラフを数秒で分析できるという謳い文句です。

LlamaIndexとは？

生成AIでアプリケーションを開発するとき、RAGなどの応用的な手法を手軽に実装するためにLangChainなどのフレームワークを活用することも多いでしょう。LLM界隈ではトップシェアのLangChainが一番有名ですが、二番手のLlamaIndexはとりわけRAG関連の機能が充実しています。

今年の5月にAmazon NeptuneのGraphRAGサポートも行われているため、今回はこのLlamaIndexを利用してみます。

実際に試してみよう！

新技術を学ぶ最良の方法は、自分で触ってみることです。理解の解像度が格段に上がります。

以下の構成を作ります。

AWS利用料が発生することをご認識ください。なるべく低スペックに抑えた手順を案内していますが、検証後は速やかに不要なリソースを削除しましょう。

AWSアカウント作成

AWSアカウントをお持ちでない方は、無料でアカウントを作成しましょう。

※今回はすべてバージニア北部リージョンを利用しますが、東京など他の主要リージョンでも概ね実施可能です。

Bedrockのモデル有効化

今回、RAGを試すうえで「回答生成用」と「埋め込み用」の2種類の生成AIモデルを利用します。事前に有効化しておきましょう。

マネジメントコンソールにサインインしたら、Amazon Bedrockのコンソールを検索して移動します。

左サイドバーを開き、「モデルアクセス」から「Enable specific models」をクリックし、以下のモデルを選択して有効化します。（所属会社や用途を聞かれるため、ざっくり回答しましょう）

• Anthropic > Claude 3.5 Sonnet
• Amazon > Titan Text Embeddings V2

Neptune Databaseを作成

マネジメントコンソールにサインインしたら、Amazon Neptuneのコンソールを検索して移動します。

左ペインの「データベース > クラスター」より、データベースの作成を実施します。

• テンプレート：開発とテスト
• クラスターストレージ設定 > 設定オプション：Neptune Standard
• インスタンスオプション > インスタンスクラス：バースト可能クラス > 'db.t4g.meduim'
• ノートブック設定
  • ノートブック名：graphrag
  • IAMロール：AWSNeptuneNotebookRole- graphrag

他はデフォルトのまま「データベースを作成」します。

グラフDBのリソースは比較的高額ですが、今回の設定で無料利用枠のインスタンスを利用することができます。サーバーレスではなくプロビジョンド版を使うのがミソです。

この初期設定により、Neptune DBクラスターと同じVPC内にSageMakerのノートブックインスタンスも作成してくれるため、お手軽にDBを操作することができます。

ノートブックでサンプルコードを実行

Neptuneコンソールの左サイドバーから「ノートブック」にアクセスし、自動で作成されているノートブックインスタンスを選択してJupyterLab（対話型のIDE。主に機械学習で使われる）を開きます。

「Notebook > Python 3」をクリックして、Python用ノートブックを新規作成します。

今回はLlamaIndex公式ドキュメントに記載されている、Amazon Neptune用のサンプルコードを実行してみます。

最初のコードブロックに以下を入力し、「▶︎」ボタンをクリックして実行します。

これは、このノートブックインスタンスに今回利用するPythonの外部ライブラリをインストールするコマンドです。

Untitled.ipynb

# Python外部ライブラリのインストール
%pip install -U boto3 llama-index-llms-bedrock llama-index-graph-stores-neptune llama-index-embeddings-bedrock llama-index-readers-file

実行すると、コードブロック単位で実行結果が表示されます。このように、少しずつ実行結果を確認しながらコードを追加していけるのが「ノートブック」の便利なところです。

Pythonライブラリの依存性エラーが赤字で表示されることがありますが、動作に支障はないためそのまま進めて大丈夫です。

同様に続けます。

次は、インストールしたライブラリをPythonコード上で使えるようにインポートします。

Untitled.ipynb

# Python外部ライブラリのインポート
from llama_index.llms.bedrock import Bedrock
from llama_index.embeddings.bedrock import BedrockEmbedding
from llama_index.core import StorageContext, SimpleDirectoryReader, KnowledgeGraphIndex, Settings
from llama_index.graph_stores.neptune import NeptuneDatabaseGraphStore
from IPython.display import Markdown, display

実行（▶︎）すると、数秒で完了します。

ちなみに実行ステータスがかなり分かりづらいのですが、画面下部のステータスバーで Python 3 | Busy もしくは Python 3 | Idle という形で、コードを処理中か否かが表示されています。

次は以下を貼り付けて実行します。回答生成用LLM（Claude）と、ベクトル変換用の埋め込みモデル（Amazon Titan）を設定しています。

Untitled.ipynb

# 回答生成用LLMと、埋め込み用モデルを設定
llm = Bedrock(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
embed_model = BedrockEmbedding(model="amazon.titan-embed-text-v2:0")

Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512

「社内文書」をグラフDBに登録

ここで、GraphRAGの検索対象となるドキュメントを準備します。

実際の社内文書の代わりに、Wikipediaの「東方永夜抄」というゲーム作品のページをPDF形式でダウンロードしておきます。（画面上部の「ツール」から実施可能）

ノートブックのルート階層に Data というディレクトリを作成し、その中にこのPDFファイルをアップロードしておきましょう。（ドラッグ＆ドロップでアップロード可能）

その後、ブラウザの別タブでNeptuneのコンソールを開き、作成したクラスターの書き込みエンドポイント名をコピーします。

ノートブックに戻り、次のコードブロックには以下を入力します。
今コピーしたエンドポイント名をダブルクォート内に貼り付けてから実行しましょう。

Untitled.ipynb

# ドキュメントとグラフDBを設定
documents = SimpleDirectoryReader("./Data").load_data()
graph_store = NeptuneDatabaseGraphStore(
    host="<Neptuneクラスターの書き込みエンドポイント名>",
    port=8182
)

次のコードブロックでは、Amazon Bedrockの生成AIモデルを呼び出します。ただし、ノートブックに実行権限が付与されていないので設定しましょう。

マネジメントコンソールからAmazon SageMakerを開き、「Notebooks」より今回利用しているノートブックインスタンスの詳細を開くと、「IAMロールARN」からIAMロール設定画面に飛ぶことができます。

このIAMロールに「ポリシーをアタッチ」で AmazonBedrockFullAccess の許可を追加しておきましょう。

その後、ノートブックに戻り、次のコードブロックで以下を実行します。
Wikipediaから抽出したテキストを実際にグラフDBへ登録する処理のため、数分かかります。

Untitled.ipynb

# インデックス作成処理（数分かかります）
storage_context = StorageContext.from_defaults(graph_store=graph_store)
index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=2,
)

個人的にGraphRAGでは埋め込みを使っているのかいないのかよく分かっていなかったのですが、前述のコードでEmbeddingモデルを指定していることからも分かるように、グラフ用のインデックス作成時に埋め込みを利用していますね。

LlamaIndexの該当コードは以下です。
https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/indices/knowledge_graph/base.py

NotebookのステータスがIdleに戻ったら、次へ進みましょう。

グラフを可視化してみる

ドキュメントをグラフDBへ登録できたら、RAGを試す前に一度、グラフを可視化してみましょう。

Neptuneコンソールの「ノートブック」より、今回のノートブックを「Graph Explorer」で開いてみます。

グラフエクスプローラーが開いたら、画面右下の「Entity」を開きます。

数ページにわたって大量の「ノード」が登録されていることが分かります。
とりあえず1画面分でいいので、画面右の「Send to Explorer」をクリックしまくって、片っ端から可視化対象に追加してみましょう。

その後、画面右上の「Open Graph Explorer」をクリックすると、可視化画面に遷移します。
追加した「ノード」が表示されており、これらをダブルクリックすると関連する「エッジ」が展開されます。ひたすら色んなノードをダブルクリックしてみましょう。

画面上部の更新ボタンを押すと、ノードとエッジを再整頓することができます。
埋め込みを使って、このようにドキュメント内の単語同士を意味にもとづいて関係づけているんですね！

いざ、GraphRAGを実行！

それでは、ノートブックに戻って実際にGraphRAGを実行してみましょう。

新しいコードブロックで以下を実行します。

Untitled.ipynb

# クエリーを実行して、結果を表示
query_engine = index.as_query_engine(
    include_text=False,
    response_mode="tree_summarize"
)
response = query_engine.query("蓬莱山輝夜が求婚されたのは誰？")
display(Markdown(f"{response}"))

出力例は以下のようになります。

蓬莱山輝夜が求婚されたのは貴族の父でした。輝夜は月の世界の重要な人物であり、高い地位にあったことがうかがえます。彼女は月人であり、月の賢者としても知られていました。また、八意永琳から教育を受け、あらゆる薬を作る能力を持っていたようです。このような背景から、貴族の父が彼女に求婚したことは、輝夜の社会的地位や能力の高さを反映していると考えられます。

正解ですね。さらに踏み込んで、藤原妹紅の父であることに言及してほしかったですが、どうもWikipediaのフォーマットを解釈する際、見出しと本文の関係性がいまいち汲み取られていないようで、思ったほどの精度が出ていない部分があります。

通常のRAGと比較してみよう！

今回のドキュメントと質問文を、GraphRAGではない通常のRAGで同様に試してみると回答できませんでした。

Bedrockのナレッジベース機能を使えば数分で簡単にRAGが試せますので、今回のGraphRAGと比較してみましょう。ユースケースによって両者の得意・不得意が如実に出てくるため面白いですよ！

エージェントでも「グラフ」が熱い！

今回のRAGに限らず、AIエージェントでもグラフを活用したツールが盛り上がっています。生成AIアプリケーションに取り組まれている方は、ぜひ早めにグラフを学んでおきたいところです。

参考文献