高度なRAG検索戦略：Auto-Merging Retrieval (自動マージングリトリーバル）

2024年5月10日 02:37

以前、私たちは文章ウィンドウ検索の高度なRAG検索戦略を探求しました。今日は、別の洗練された検索戦略である自動マージングについて詳しく見ていきましょう。文章ウィンドウ検索よりもやや複雑ですが、以下の説明によってその原理を理解するのに役立ちます。また、LlamaIndexを使用して自動マージング検索を構築する方法と、最終的にTrulensを使用してその検索効果を評価し、以前の検索戦略と比較します。

自動マージング検索の紹介

自動マージング検索は、主に異なるブロックサイズのノードに文書を分割し、親ノードと子ノードを含みます。検索プロセス中に、類似度の高いリーフノードが特定されます。親ノード内で複数の子ノードが検索された場合、その親ノードは自動的にマージされます。最終的に、親ノードからのすべての文書がLLM（Large Language Model）へのコンテキストとして送信されます。以下は、自動マージング検索プロセスのイラストです。

自動マージ検索は、LlamaIndex内の高度な検索機能であり、ドキュメントの分割とマージという2つの主要なプロセスが関与しています。コード例を通じて基本原則を探ってみましょう。

ドキュメントの分割

自動マージ検索を構築するためには、まずHierarchicalNodeParserドキュメントパーサーを作成する必要があります。

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import HierarchicalNodeParser
documents = SimpleDirectoryReader("./data").load_data()
node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=[2048, 512, 128])
nodes = node_parser.get_nodes_from_documents(documents)

ますます、私たちはdataディレクトリからドキュメントをロードします。
次に、HierarchicalNodeParserドキュメントパーサーを作成し、chunk\_sizesを[2048, 512, 128]に設定します。
その後、ドキュメントパーサーを使用してドキュメントをノードに解析します。

HierarchicalNodeParserパーサーのchunk\_sizesパラメータはデフォルトで[2048, 512, 128]に設定されており、これはドキュメントが3つのレベルに分割されることを示しています：最初のレベルには2048、2番目には512、3番目には128があります。レベルを少なくしたり多くしたりすることができます。たとえば、2つのレベル[1024, 128]や4つのレベル[2048, 1024, 512, 128]などです。より小さな分割は検索の精度を高める一方で、マージの可能性を減らすことがあります。調整は評価結果に基づいて行う必要があります。

ルートノードとリーフノードの識別

LlamaIndexには、異なるレベルのノードを特定するのに役立つユーティリティ関数が用意されています。ルートノードとリーフノードを見つける方法を見てみましょう。

from llama_index.core.node_parser import get_leaf_nodes, get_root_nodes
print(f"total len: {len(nodes)}")
root_nodes = get_root_nodes(nodes)
print(f"root len: {len(root_nodes)}")
leaf_nodes = get_leaf_nodes(nodes)
print(f"leaf len: {len(leaf_nodes)}")
# Results display
total len: 66
root len: 4
leaf len: 52

メソッドget\_leaf\_nodesとget\_root\_nodesは、どちらもノードのリストを入力として受け取ります。
合計66個のノード、4つのルート、52枚の葉を観察します。
ルートと葉のノードの合計（56）は合計（66）と一致せず、中間レベルのノード（66–56）が10個あることを示しています。
ドキュメントの階層が2レベルの場合、ルートと葉のノードの合計は合計ノード数と一致します。

異なるレベルのノードの識別

別のユーティリティ関数get\_deeper\_nodesを使用して、推論を検証しましょう。

from llama_index.core.node_parser import get_deeper_nodes
deep0_nodes = get_deeper_nodes(nodes, depth=0)
deep1_nodes = get_deeper_nodes(nodes, depth=1)
deep2_nodes = get_deeper_nodes(nodes, depth=2)
print(f"deep0 len: {len(deep0_nodes)}")
print(f"deep1 len: {len(deep1_nodes)}")
print(f"deep2 len: {len(deep2_nodes)}")
# Results display
deep0 len: 4
deep1 len: 10
deep2 len: 52

get_deeper_nodesメソッドの最初のパラメータはノードリストであり、2番目のパラメータはクエリするレベルです。ここで、0は最初のレベル、つまりルートノードを表します。

これにより、deep0ノードの数は4であり、ルートノードと同等であり、deep2ノードの数は52であり、葉ノードと同等です。そして、deep1ノードは10個の中間レベルにあり、先ほどの推論と一致しています。

子ノードの識別

LlamaIndexには、ノードの子ノードを特定するget_child_nodes関数も提供されています。

from llama_index.core.node_parser import get_child_nodes
middle_nodes = get_child_nodes(root_nodes, all_nodes=nodes)
leaf_nodes = get_child_nodes(middle_nodes, all_nodes=nodes)
print(f"middle len: {len(middle_nodes)}")
print(f"leaf len: {len(leaf_nodes)}")
# Results display
middle len: 10
leaf len: 52

get\_child\_nodesメソッドの最初のパラメータは子ノードを取得するためのノードリストであり、2番目のパラメータはすべてのノードです。
ここでは、まずルートノードのすべての子ノードを取得し、10個の中間レベルのノードが得られます。
次に、これらの中間ノードのすべての子ノードを特定し、52個の葉ノードが得られます。

もちろん、最初のルートノードの子ノードのように、特定のノードの子ノードも特定できます。

root0_child_nodes = get_child_nodes(root_nodes[0], all_nodes=nodes)
print(f"root0 child len: {len(root0_child_nodes)}")
# Results display
root0 child len: 2

これは、最初のルートノードが2つの子ノードを持ち、それらが中間レベルのノードであることを示しています。

ノードのドキュメントコンテンツ

各親ノードのドキュメントコンテンツには、すべての子ノードのドキュメントコンテンツが含まれています。

print(f"deep1[0] node: {deep1_nodes[0].text}")
child = get_child_nodes([deep1_nodes[0]], all_nodes=nodes)
print(f"child[0] node of deep1[0]: {child[0].text}")
# Results display
deep1[0] node: In the Eastern European country of Sokovia, the Avengers-Tony Stark, Thor, Bruce Banner, Steve Rogers, Natasha Romanoff, and Clint Barton-raid a Hydra facility commanded by Baron Wolfgang von Strucker, who has experimented on humans using the scepter previously wielded by Loki. They meet two of Strucker's test subjects-twins Pietro (who has superhuman speed) and Wanda Maximoff (who has telepathic and telekinetic abilities)-and apprehend Strucker, while Stark retrieves Loki's scepter.
Stark and Banner discover an artificial intelligence within the scepter's gem, and secretly decide to use it to complete Stark's "Ultron" global defense program. The unexpectedly sentient Ultron, believing he must eradicate humanity to save Earth, eliminates Stark's A.I. J.A.R.V.I.S. and attacks the Avengers at their headquarters. Escaping with the scepter, Ultron uses the resources in Strucker's Sokovia base to upgrade his rudimentary body and build an army of robot drones. Having killed Strucker, he recruits the Maximoffs, who hold Stark responsible for their parents' deaths by his company's weapons, and goes to the base of arms dealer Ulysses Klaue in Johannesburg to get vibranium. The Avengers attack Ultron and the Maximoffs, but Wanda subdues them with haunting visions, causing Banner to turn into the Hulk and rampage until Stark stops him with his anti-Hulk armor.[a]
A worldwide backlash over the resulting destruction, and the fears Wanda's hallucinations incited, send the team into hiding at Barton's farmhouse. Thor departs to consult with Dr. Erik Selvig on the apocalyptic future he saw in his hallucination, while Nick Fury arrives and encourages the team to form a plan to stop Ultron. In Seoul, Ultron uses Loki's scepter to enslave the team's friend Helen Cho. They use her synthetic-tissue technology, vibranium, and the scepter's gem to craft a new body. As Ultron uploads himself into the body, Wanda is able to read his mind; discovering his plan for human extinction, the Maximoffs turn against Ultron.
child[0] node of deep1[0]: In the Eastern European country of Sokovia, the Avengers-Tony Stark, Thor, Bruce Banner, Steve Rogers, Natasha Romanoff, and Clint Barton-raid a Hydra facility commanded by Baron Wolfgang von Strucker, who has experimented on humans using the scepter previously wielded by Loki.

以下は記事の翻訳です：

まず、最初の中間レベルノードのドキュメントコンテンツを印刷します。
次に、この中間ノードの最初の子ノードのドキュメントコンテンツを取得して印刷します。
親ノードのドキュメントコンテンツには、子ノードのものが含まれていることが明らかです。

ドキュメントマージング

ドキュメントマージングは自動マージング検索の重要なコンポーネントです。ドキュメントマージングの結果は、LLMに提出されるコンテキストコンテンツを決定し、最終的な生成結果に影響を与えます。
最初に、自動マージング検索はクエリに基づいてすべてのリーフノードを検索し、検索の高い精度を確保します。自動マージング検索では、simple\_ratio\_threshというパラメータがあり、デフォルト値は0.5です。このパラメータは自動ドキュメントマージングの閾値を設定します。親ノード内の取得された子ノードの比率がこの閾値を下回る場合、自動マージング機能はアクティブにならず、LLMに提出されるコンテキストには取得されたリーフノードのみが含まれます。逆に、比率がこの閾値を上回る場合、ドキュメントは自動的にマージされ、LLMに提出されるコンテキストにはこの親ノードのコンテンツが含まれます。
例えば、親ノードに4つの子ノードがあり、取得された子ノードが1つだけの場合、取得された子ノードの比率は0.25 (1/4) であり、0.5の閾値を下回ります。したがって、自動マージング機能はアクティブにならず、LLMに提出されるコンテキストには取得された子ノードのみが含まれます。

親ノードが4つの子ノードを持ち、そのうち3つが取得された場合、取得された子ノードの比率は0.75（3/4）となり、0.5の閾値を超えます。したがって、自動マージ機能が作動し、LLMに提出されるコンテキストは親ノードの内容となります。

さらに、自動マージ機能は繰り返しのプロセスであり、自動マージは最下位のノードから開始し、最終的にすべてのマージされたドキュメントを取得するために、最上位のノードまでマージを続けます。繰り返しの回数は、ドキュメントパーサーによるドキュメントの分割レベルに依存します。たとえば、chunk\_sizes が [2048, 512, 128] の場合、3つのレベルのドキュメント分割が発生します。各レベルで自動マージがトリガーされる場合、2回マージされます。

自動マージの利用

実践的なRAGプロジェクトでの自動マージ検索の利用をさらに探求し、Wikipediaのアベンジャーズ映画のプロットをドキュメントデータとして使用してテストします。

自動マージ検索の例

LlamaIndexを使用して自動マージ検索を構築する方法は次のとおりです：

from llama_index.core.node_parser import (
    HierarchicalNodeParser,
    get_leaf_nodes,
)
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import SimpleDirectoryReader
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.settings import Settings
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=[2048, 512, 128])
documents = SimpleDirectoryReader("./data").load_data()
nodes = node_parser.get_nodes_from_documents(documents)
leaf_nodes = get_leaf_nodes(nodes)
llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
embed_model = OpenAIEmbedding()
Settings.llm = llm
Settings.embed_model = embed_model
Settings.node_parser = node_parser
docstore = SimpleDocumentStore()
docstore.add_documents(nodes)
storage_context = StorageContext.from_defaults(docstore=docstore)
base_index = VectorStoreIndex(leaf_nodes, storage_context=storage_context)
base_retriever = base_index.as_retriever(similarity_top_k=12)
retriever = AutoMergingRetriever(
    base_retriever,
    storage_context,
    simple_ratio_thresh=0.3,
    verbose=True,
)
auto_merging_engine = RetrieverQueryEngine.from_args(retriever)

最初に、以前に紹介したように、HierarchicalNodeParser ドキュメントパーサーを定義してドキュメントを解析します。
次に、OpenAIのLLMと埋め込みモデルを使用して回答生成とベクトル作成を行います。
その後、すべてのノードnodesを保存するためのstorage\_contextを作成します。その後の自動マージング検索では、リーフノードに基づいて関連する親ノードを見つけるため、すべてのノードを保存する必要があります。
その後、基本的な検索base\_indexを構築し、質問に基づいてすべてのリーフノードleaf\_nodesを検索し、similarity\_top\_kに最も一致するノードを特定します。ここでは、最も一致するリーフノード12個を取得します。
次に、自動マージング検索AutoMergingRetrieverを構築します。これは基本検索の結果に基づいてマージします。ここでは、simple\_ratio\_threshを0.3に設定し、この閾値を超える子ノードの検索比率を持つノードは自動的にマージされます。verboseパラメータはTrueに設定されており、マージプロセスの出力を示します。
最後に、RetrieverQueryEngineを使用して検索エンジンを作成します。

これで、この検索エンジンを使用して質問に答えることができます。

question = "Who created Ultron in the Avengers?"
response = auto_merging_engine.query(question)
print(f"response: {str(response)}")
print(f"nodes len: {len(response.source_nodes)}")
# Display results
> Merging 5 nodes into parent node.
> Parent node id: f07e7918-3bf6-4b03-83a8-cb3b5370c856.
> Parent node text: In the Eastern European country of Sokovia, the Avengers-Tony Stark, Thor, Bruce Banner, Steve Ro...
> Merging 4 nodes into parent node.
> Parent node id: f0fbd689-f92a-41f5-9e15-ec0e6433e680.
> Parent node text: Rogers, Romanoff, and Barton fight Ultron and retrieve the synthetic body, but Ultron captures Ro...
> Merging 2 nodes into parent node.
> Parent node id: 99a2ca9b-6cf4-4165-a8d0-c97853fa13e1.
> Parent node text: In the Eastern European country of Sokovia, the Avengers-Tony Stark, Thor, Bruce Banner, Steve Ro...
response: Tony Stark and Bruce Banner
nodes len: 4

自動マージの前に、私たちのベース検索は最も一致する12個のリーフノードを取得しました。出力から、これらの12個のノードが3つのマージ操作を経て、最終的に4つのノードになったことが観察されました。これらのノードには、マージされたリーフノードと親ノードの両方が含まれています。

検索効果の比較

その後、自動マージ検索の効果を評価するためにTrulensを使用しました。

tru.reset_database()
rag_evaluate(base_engine, "base_evaluation")
rag_evaluate(sentence_window_engine, "sentence_window_evaluation")
rag_evaluate(sentence_window_engine, "auto_merging_evaluation")
Tru().run_dashboard()

rag\_evaluateの詳細については、私の以前の記事を参照してください。これは、Trulensのgroundedness、qa\_relevance、qs\_relevanceを主に使用して、RAGの検索結果を評価します。以前の基本的な検索と文ウィンドウ検索の評価を保持し、自動マージ検索の評価を追加しました。コードを実行した後、ブラウザのTrulensで評価結果を確認できます。

評価結果では、自動マージ検索が他の2種類の検索よりも優れていることがわかります。ただし、これは自動マージ検索が常に他の検索を上回ることを意味するわけではありません。具体的な評価結果は、元の入力文書、検索パラメータの設定などに依存します。要するに、具体的な評価結果は実際の状況に基づいて評価されるべきです。

結論

自動マージ検索は高度なRAG検索の手法です。文書の分割と結合のアイデアがこの手法の主な特徴です。この記事では、自動マージ検索の原則と実装方法を紹介し、Trulensを使用してその効果を評価しました。自動マージ検索をよりよく理解し、活用するためのお手伝いができればと思います。

データ元：

https://generativeai.pub/advanced-rag-retrieval-strategies-auto-merging-retrieval-dc3f869654c4

この記事が気に入ったらサポートをしてみませんか？