2024年12月4日、Amazon Bedrock Knowledge basesがGraphRAGをサポートするようになったというアップデートがありました！

ただ、以前の検証で「このGraphRAGあんまり精度高くなくね…？」となりました。
まだプレビュー版ということもあると思いますが、現状だとBaseline RAGとそれほど変わらない精度でした。
参考として以下の記事に結果をまとめてあります。

ただ、これとは別の方法でもGraphRAGをAWS上で構築できるようです。
今回はそれを試してみて、ナレッジベース上から作成するGraphRAGと精度を比較することが目的です。

実際に構築してみる

GraphRAGを構築するにあたって、今回参考にさせていただいた記事は以下です。

今回はナレッジベースの機能としてではなく、LlamaIndexとAmazon Neptuneを用いてGraphRAGを構築します。

Amazon Neptuneとは、フルマネージド型のグラフデータベースです。
ここでは、GraphRAGにおけるナレッジグラフを作成するために用います。

LlamaIndexとは、生成AIアプリ開発で利用できるフレームワークの1つです。
とりわけRAG関連の機能が充実しており、公式としてAWS上でGraphRAGを構築する方法が記載されているので、こちらを利用します。

使用するAWSリソース

• Amazon Bedrock
  • Claude 3.5 Sonnet(回答生成用)
  • Amazon Titan Embeddings G1 - Text(埋め込み用)
• Amazon Neptune
  • データベース(プロビジョンド)
  • Jupyter Notebook(実態はSageMaker AI)

Amazon Bedrockのモデルアクセス有効化

Bedrockコンソールを開き、以下2つのモデルアクセスを有効化します。
・Amazon Titan Embeddings G1 - Text

・Claude 3.5 Sonnet

Amazon Neptuneの構築

Neptuneコンソールにおいて、メニュー一覧から「データベース」配下にある「クラスター」を選択し、「データベースを作成」をクリックします。

作成画面では、以下のパラメータを選択します。記載していないものはデフォルトでOKです。

• インスタンスタイプ：Provisioned
• テンプレート：開発とテスト
• クラスターストレージ設定：Neptune Standard
• インスタンスクラス：バースト可能クラス -> db.t4g.medium
• VPCセキュリティグループ：新規作成 (後で少し書き加えます)
  • セキュリティグループ名：NeptuneSecurityGroup
• ノートブックを作成：有効化
  • インターネットアクセス：Amazon SageMakerを介した直接アクセス

Neptuneにはサーバーレス版もありますが、サーバーレスとは名ばかりの従量課金制です。最低利用料金がかかってくるスタイルですね。

それよりも、プロビジョンド版で無料利用枠を使ったほうが安く済むそうなのでこちらを使用します。

また、ここでノートブックというものも作成します。具体的にはJupyter Notebookですね。

Jupyter Notebookとは、ライブコードとナレーションテキストを含む、完全マネージド型のインタラクティブな開発環境です。
主に機械学習・データ解析などのデータサイエンスに利用されます。

Neptune DBクラスターと同じVPC内にノートブックが作成されるので、ノートブック経由で簡単にNeptuneを操作することができます。

ちなみに、Jupyter Notebookの裏側はSageMaker AIのノートブックインスタンスです。
SageMaker AIのコンソール上でもノートブックが確認できます。

また、1点ネットワーク周りで追加設定が必要です。
DB作成時、一緒に作成されたセキュリティグループを以下のように設定してください。

• インバウンドルール：
  • タイプ：カスタムTCP
  • プロトコル：TCP
  • ポート範囲：8182
  • ソース：NeptuneSecurityGroup
• アウトバウンドルール
  →デフォルトで設定されている0.0.0.0/0宛のルールはそのままでOKです
  • タイプ：カスタムTCP
  • プロトコル：TCP
  • ポート範囲：8182
  • 送信先：NeptuneSecurityGroup

このセキュリティグループはNeptuneインスタンスだけでなく、Jupyter Notebookにもアタッチされるので、このような設定にします。

(ネットワーク周りでトラブりまくった結果この設定に行き着きましたが、他にも良い設定があるかもしれません。)

ナレッジグラフを作成する

ここからはノートブックからNeptune内に接続し、ナレッジグラフを作成していきます。

作成したノートブックを選択して、「JupyterLabを開く」を選択します。

その後、「Launcher」画面で「Python 3」を選択します。

その後、以下コマンドをコードブロックに入力し、実行していきます。
実行するには、画面上部の「▶」を押すか、Shift+Enterを押します。

1.必要なパッケージをインストールする

ターミナル上でpip installするのと同じイメージです。

%pip install boto3
%pip install llama-index-llms-bedrock
%pip install llama-index-embeddings-bedrock
%pip install llama-index-graph-stores-neptune
%pip install llama-index

実行後は、画面下部に注目してください。Idleになったら次へ進むようにしましょう。

• Python 3 | Busy -> 処理実行中
• Python 3 | Idle -> 処理完了済み
  

2.インストールしたライブラリをPythonコード上で使用できるようにインポートする

コードでimport文を書くのと同じイメージです。

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.core.query_engine import KnowledgeGraphQueryEngine
from llama_index.graph_stores.neptune import NeptuneDatabaseGraphStore
from IPython.display import Markdown, display

3.変数を登録する

デプロイ時に使用する変数を登録します。

region_name = "<NotebookやNeptuneを作成したリージョン>"
llmodel = "anthropic.claude-3-5-sonnet-20240620-v1:0"
embed_model = "amazon.titan-embed-text-v1"
neptune_endpoint = "<Neptune DBクラスターの書き込みエンドポイント>"

4.指定したAmazon Bedrockを使用するようにLLMを設定する

上記の3.で指定したBedrockの基盤モデルを使用するように設定します。

llm = Bedrock(
  model=llmodel,
  region_name=region_name,
)
embed_model = BedrockEmbedding(model=embed_model)
Settings.llm = llm
Settings.chunk_size = 512
Settings.embed_model = embed_model

Settings.chunk_sizeは、テキストデータを処理する際の分割サイズを指定するパラメータです。
後ほどまた登場します。

5.検索対象ドキュメントをNeptuneに登録

まずは検索対象ドキュメントをNotebook内に格納します。
ノートブック上で上部のフォルダ作成ボタンを押し、ルート直下にDataフォルダを作成します。

作成したDataフォルダ内に、検索対象ドキュメントを格納します。
ここでは競走馬であるドウデュースのドキュメントを使用します。

そして以下のコマンドを実行し、ドキュメントを読み込んで処理できるようにします。

documents = SimpleDirectoryReader("./Data").load_data()
print(documents)

printコマンドで、どのように読み込まれているかが確認できます。
特に日本語PDFの場合は別の方法で読み込まないと文字化けすることもあるらしい？ので一度確認することをおすすめします。

確認できたら、そのドキュメントを元にナレッジグラフを作成していきます。

graph_store = NeptuneDatabaseGraphStore(host=neptune_endpoint, port=8182)
storage_context = StorageContext.from_defaults(graph_store=graph_store)
index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=2,
)

1行目でNeptuneデータベースへの接続設定を行っています。

2行目でグラフデータを格納するためのストレージコンテキストを初期化し、作成したNeptune接続設定を用いるように指定しています。

3行目以降で、ナレッジグラフを作成していきます。具体的な設定は以下の通りです。
・documents：入力となるドキュメント(先ほどDataフォルダ配下に格納したもの)
・storage_context：作成したストレージコンテキスト
・max_triplets_per_chunk=2：各テキストチャンクから抽出する関係(トリプレット)の最大数を2とする

実行すると、恐らく数分待つと思います。
ステータスがIdleになれば、ナレッジグラフ作成完了です！

グラフを可視化してみる

Neptuneコンソールのノートブックから「Graph Explorer」を開きます。

Entityをクリックします。

すると、ドキュメントから抽出されたノードが表示されます。画面右側にある「Send to Explorer」を片っ端からクリックしましょう。
これによって、グラフの中で確認できるようになります。

登録できたら、右上の「Open Graph Explorer」をクリック。

すると、こんな風に大量のノードが出てきます。ただ、これだと何がどうなってるのか全くわかりませんね。

そこで、このノードをダブルクリックしましょう。そのノードに紐づくエッジが出現します。
各ノードの上に書いてある各数字が、そのノードの持つエッジ数ですね。

すべてのノードをダブルクリックしたら、画面上部の更新マークを押します。
すると、ノードとエッジの位置が整理整頓されます。

拡大してみると、エッジがどういった関係性を持っているかが表示されます。
数が多いとごちゃっとして見にくいですが…
なにはともあれ、これでナレッジグラフが作成できました。

クエリを投げてみる

ノートブックに戻って、以下コマンドからクエリを投げてみましょう。

response = index.as_query_engine().query(
"<好きな質問を入力>",
)
display(Markdown(f"<b>{response}</b>"))

↓

ふむふむ。まぁ正解でしょうかね。ただ、本当は合計3回戦ったので、全部のレース戦績を出してほしいところです。
ということでもう1つ質問。

あれ…上手く出力されませんでした。

追加でもう1つ質問。

こちらも概ね正しいのですが、すべてではないですね。全部で5レース走っているので、3/5だけきちんと出力されています。

上記3つの質問からわかったことは、2022年で時空が止まっており、2023年以降のデータが入っていないこと。

ということで、もしかしたらドキュメントの埋め込みが一部で止まってしまっているのかもしれません。
その原因になりうるのは2箇所。

1. Settings.chunk_sizeの値が小さい
2. max_triplets_per_chunkの値が小さい

Settings.chunk_sizeは各チャンクのサイズなのであまり関係ないだろうと思いつつも、もう少し大きくても良いと判断して、Embeddingで使用したTitanの上限である1536まで増やし、コマンドを実行します。

# Settings.chunk_sizeを変更する
llm = Bedrock(
model=llmodel,
region_name=region_name,
)
embed_model = BedrockEmbedding(model=embed_model)
Settings.llm = llm
Settings.chunk_size = 1536 #512から1536へ変更
Settings.embed_model = embed_model

そして、max_triplets_per_chunkも変更します。ここでは2から10に増やします。
以下のコマンドを実行。

# 同じドキュメントと同じストレージコンテキストを使用して、新しいパラメータでインデックスを再作成
storage_context = StorageContext.from_defaults(graph_store=graph_store)
index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=10, #2から10に変更
)

# 以下のテキストが1チャンクとする:
text = """
AWSはクラウドサービスを提供しています。
EC2はAWSのコンピューティングサービスです。
S3はAWSのストレージサービスです。
Lambda はサーバーレスサービスです。
"""

# max_triplets_per_chunk=2 の場合、抽出されるトリプレットは:
# 1: (AWS, 提供する, クラウドサービス)
# 2: (EC2, 部分である, AWSサービス)
# それ以下のトリプレットは抽出されない

そして再度同じ質問をすると…？

ｷﾀ━━━━(ﾟ∀ﾟ)━━━━!!
これこれ！！これですよこれ！！僕が求めていたものは！！！

もう1つの方はニエル賞が消えましたが、4/5レース分の結果が出力されているので、さっきのよりは全然良いです！
(多分Settings.chunk_sizeも増やしたためだと思いますが)

ということで、チャンク数やトリプレットはチューニングするのが重要そうです！

ナレッジベースから作れるGraphRAGと比較する

以前作成したナレッジベースのGraphRAGでは、上記の質問に対してまともに答えられませんでした。

単純に私の検証における精度のみを比較すると、LlamaIndexを用いてGraphRAGを作成した方が良さそうです。

ただ、ナレッジベースの方が圧倒的に構築は楽ですし、手間もかかりません。
また、出力のソースがどこだったのかわかるので、かなりデバッグやチューニングがしやすそうです。

何よりre:Inventでアップデートが大量に来ていることからも、今後もより一層機能が充実していきそうです。

ということでまとめるとこんな感じになりそうです(ふくち調べ)。

※間違っている可能性大いにあります！
※「そうなんだな〜自分も試してみるか〜」くらいな感じでお願いしますmm

まとめ

ナレッジベース上で作れるGraphRAGの精度が、LlamaIndexで作成するGraphRAGくらい良いものになりますように！！

今後の伸びしろは抜群なので、GAに期待！

おまけ

最後に、出走が叶わなくなってしまった2024年有馬記念、ドウデュースの引退レース結果を妄想してもらいました。

「最終コーナーを回ったところで武豊騎手が外に持ち出し、一気に加速」
この光景が目に浮かぶようです。

最後にはBedrockも言葉を紡げなくなるほど感動して泣いてしまいましたね()。

たくさんの夢をありがとう、ドウデュース。