Amazon Neputune が GraphRAG Toolkit をリリースしました
こんにちは!クラウド事業本部コンサルティング部のたかくに(@takakuni_)です。
少し前ですが Amazon Neptune チームが OSS で新しく GraphRAG Toolkit をリリースしました。
GraphRAG Toolkit
GraphRAG Toolkit は自前で GraphRAG アプリケーションを構築するために利用する Python フレームワークです。
非構造化データから、グラフの構築(インデックス化)を自動化したり、ユーザーからの質問に回答するためのクエリ処理をサポートします。
データコネクタ、メタデータ抽出、変換には LlamaIndex コンポーネントを利用して処理を行います。
Amazon Bedrock Knowledge bases では、ベクターストア/グラフストアの両方を Amazon Neptune Analytics で賄いますが、 GraphRAG Toolkit では次のように使い分けることができます。[1]
- グラフストア
- Amazon Neptune Analytics
- Amazon Neptune Database
- ベクターストア
- Amazon Neptune Analytics
- Amazon OpenSearch Serverless
やってみる
それでは実際に GraphRAG Toolkit を利用して GraphRAG を作ってみます。今回は examples が用意されているためこちらを利用します。リージョンはオレゴンリージョンを使いました。
CloudFromation のデプロイ
CloudFormation テンプレートが用意されています。こちらを使って、ベクターストア/グラフストア/NoteBook インスタンスを作成します。(おおよそ 15 分程度で完成します。)
完成したら Notebook インスタンスから Jupiter を開きます。
データの抽出と構築
こちらのノートブックを利用します。
GraphRAG Toolkit のインデックス化には、Extract(抽出)と Build(構築)の 2 段階のフェーズがあります。
Extract
Extract では以下の抽出処理を行います。(かなり本格的ですね)
- ソースドキュメントのチャンク化
- LLM を利用して
- 複雑な文章をよりシンプルな文章に分割
- 代名詞を特定の名前に置き換え
- 可能な場合は頭字語を置き換え
- LLM を利用して、エンティティ、関係、トピック、ステートメント、ファクトを抽出
チャンク化が済んでいればステップ 1 は省略可能。ステップ 2 も LLM のコストカットとパフォーマンス低下をトレードオフにして良ければ省略可能です。ステップ 3 は必須です。
Build
続いて Build です。 Extract フェーズで行われたデータでインデックスを作成します。抽出されたデータはソース、チャンク、トピック、ステートメント、ファクトに分割され、チャンクノードとステートメントノードがインデックス化の対象になります。
extract_and_build
前置きが長くなりましたが、extract_and_build()
は上記の Extract と Build を一括に処理する関数です。実際に実行してみましょう。
モデルやスレッド数を変更したい場合は、以下の GraphRAGConfig
を参考に変更します。
たとえば抽出で利用する LLM を変更する場合は次のように記載します。
%reload_ext dotenv
%dotenv
import os
from graphrag_toolkit import LexicalGraphIndex, GraphRAGConfig
from graphrag_toolkit.storage import GraphStoreFactory
from graphrag_toolkit.storage import VectorStoreFactory
from llama_index.readers.web import SimpleWebPageReader
import nest_asyncio
nest_asyncio.apply()
+ GraphRAGConfig.extraction_llm = 'anthropic.claude-3-haiku-20240307-v1:0'
graph_store = GraphStoreFactory.for_graph_store(os.environ['GRAPH_STORE'])
vector_store = VectorStoreFactory.for_vector_store(os.environ['VECTOR_STORE'])
graph_index = LexicalGraphIndex(
graph_store,
vector_store
)
doc_urls = [
'https://docs.aws.amazon.com/neptune/latest/userguide/intro.html',
'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html',
'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-features.html',
'https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-vs-neptune-database.html'
]
docs = SimpleWebPageReader(
html_to_text=True,
metadata_fn=lambda url:{'url': url}
).load_data(doc_urls)
graph_index.extract_and_build(docs, show_progress=True)
print('Complete')
抽出と構築の分離
先ほど extract_and_build
でデータの抽出を構築を一括に行っていましたが、分離して実行することもできます。
たとえば、開発環境/本番環境で同じデータを利用したいケースでは、抽出を 1 回、構築を 2 回と言った形で抽出の LLM のコストをカットできます。
Extract
Extract を分割する場合、S3BasedChunks
と FileBasedChunks
がサポートされています。
今回は、file_based_chunks
を利用したため、ファイルパスに設定した './extracted/'
に、チャンク化されたデータが保管されていますね。
チェックポイント
GraphRAG Toolkit は、チェックポイントの仕組みが設けられています。Extract と Build の両方の処理をサポートしており、非常に便利ですね。
応用編
最後に応用編です。LexicalGraphIndex は抽象化されているため、もっと細かく設定したい場合は ExtractionPipeline
や BuildPipeline
を利用できます。
クエリ
最後にクエリの実行です。
GraphRAG Toolkit は次のリトリーバーをサポートします。各リトリーバーは、複数の検索手法を利用して検索を行います。2 つのデータベースを最大限活用して、本格的な RAG を実装しています。
- TraversalBasedRetriever
- EntityBasedSearch
- クエリからキーワードを抽出し、キーワードに基づいてグラフ内のエンティティを検索する検索手法
- ChunkBasedSearch
- ベクトル類似性検索を使用してチャンクを検索する検索手法
- EntityBasedSearch
- SemanticGuidedRetriever
- StatementCosineSimilaritySearch
- クエリとステートメントのコサイン類似度を計算して取得する検索手法
- KeywordRankingSearch
- クエリからキーワードと同義語を抽出し、全文検索を行う検索手法
- キーワードの一致数でランキング付けする
- クエリからキーワードと同義語を抽出し、全文検索を行う検索手法
- SemanticBeamGraphSearch
- クエリからキーワードを抽出し、コサイン類似度検索の結果から元のクエリと関連性の高い文章を再ランク付けする手法
- StatementCosineSimilaritySearch
TraversalBasedRetriever
### 質問文
What are the differences between Neptune Database and Neptune Analytics?
### 回答
Neptune Database and Neptune Analytics are complementary services with different purposes:
Neptune Database is a fully managed graph database service designed for high performance and availability. It is optimized for graph database workloads that require scaling to handle up to 100,000 queries per second. Neptune Database provides features like multi-AZ high availability, multi-region deployments, and is suitable for applications like fraud detection, social networking, and customer 360 views. [https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-vs-neptune-database.html, https://docs.aws.amazon.com/neptune/latest/userguide/intro.html]
On the other hand, Neptune Analytics is a memory-optimized graph analytics engine focused on quickly analyzing large graph datasets. It can load data from Neptune Database snapshots, Amazon S3, or other sources into memory for low-latency analytical queries and processing using graph algorithms. Neptune Analytics automatically provisions compute resources based on the graph size and is ideal for data science, investigative, and exploratory workloads that require fast iteration and insights on graph data. [https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html, https://docs.aws.amazon.com/neptune-analytics/latest/userguide/neptune-analytics-vs-neptune-database.html]
In summary, while Neptune Database is an operational graph database, Neptune Analytics is an analytics engine optimized for deriving insights and trends from large graph datasets using in-memory processing and graph algorithms.
SemanticGuidedRetriever
### 質問
What are the differences between Neptune Database and Neptune Analytics?
### 回答
Neptune Database and Neptune Analytics are two related but distinct services offered by AWS:
Neptune Database is a fully managed graph database service designed for high performance and scalability. It is optimized for handling highly connected datasets and transactional, operational workloads that require low latency queries and updates. [source_3]
On the other hand, Neptune Analytics is an analytics database engine focused on quickly analyzing large graph datasets stored in data lakes or existing graph databases like Neptune Database. It loads the graph data entirely into memory for low-latency analytic queries and processing using popular graph algorithms. Neptune Analytics is ideal for data science, investigative, and exploratory workloads that require fast iteration and insights over massive graph datasets. [source_1, source_2]
Some key differences are:
1. Neptune Database is a transactional database optimized for operational workloads, while Neptune Analytics is an analytics engine for processing large graphs in-memory.
2. Neptune Database provides multi-AZ high availability and multi-region deployments, whereas Neptune Analytics does not have those capabilities. [source_1]
3. Neptune Analytics complements Neptune Database by allowing you to run complex analytics on graph data without impacting the transactional database. [source_2]
4. Neptune Analytics can load data from Neptune Database snapshots or S3, analyze it quickly using algorithms, and discard the in-memory data when done. [source_2, source_4]
5. Neptune Analytics automatically provisions compute resources based on the graph size, while Neptune Database is provisioned separately. [source_4]
In summary, use Neptune Database for operational graph workloads requiring high availability and scalability, and Neptune Analytics for exploratory analytics over massive graphs using in-memory processing.
まとめ
以上、「Amazon Neptune チームが OSS として新しく GraphRAG Toolkit をリリースしました。」でした。
かなりハイレベルな GraphRAG なので、素晴らしさが伝えきれてない気がしますが、カスタムで RAG を作るのであれば、選択肢の 1 つに挙がると思いました。AWS Blog も出ていますので、合わせてご覧いただけますと幸いです。
クラウド事業本部コンサルティング部のたかくに(@takakuni_)でした!
もちろん、追加のインフラコストは発生します。たとえば OpenSearch Serverless の場合、拡張機能を利用できるなど、得意な部分でさらに精度を伸ばしたいケースで使い分けが生きてくるのではないかと思います。(シンプルにクエリ速度も変わってくるとは思います。) ↩︎