Advanced RAG Retrieval Strategy: Query Rewriting
Introduction to several query rewriting strategies in advanced RAG retrieval
RAG（Retrieval Augmented Generation）アプリケーションでは、ドキュメントの検索が高品質な回答を保証するために重要です。しかし、ドキュメントの検索だけでなく、ユーザーのクエリの最適化も同じくらい重要です。時には、ユーザーのクエリが明確でないか、具体的でないかもしれず、検索精度を向上させるためのクエリの書き換えが必要になることがあります。今日は、RAGアプリケーションにおける一部のクエリ書き換え戦略と、それらを実際のプロジェクトで使用する方法について紹介します。
サブクエリの質問
サブクエリ戦略、より小さな質問を作成するための手段です。これは、答えを見つける過程で大きな質問に関連するより小さな質問に分解して作成することです。これにより、主要な質問をよりよく理解するのが可能になります。この小さな質問は、より詳細であり、システムが主要な質問をより徹底的に理解するのに役立ちます。これにより、クエリ検索の精度が向上し、正しい答えを提供するのに役立ちます。
- まず、サブクエリ戦略は、LLM（Large Language Model）を使用してユーザーのクエリから複数のサブクエリを生成します。
- 次に、各サブクエリはRAGプロセスを経て、それぞれの答えを得ます（検索生成）。
- 最後に、すべてのサブクエリの答えが統合され、最終的な答えが得られます。
コード例
サブクエリの生成／実行はLlamaIndexに実装されています。サブクエリの生成／実行の効果を調査する前に、複雑な質問に対する通常のRAG検索のパフォーマンスを見てみましょう:
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
question = "Are Harley Quinn and Thanos righteous characters in the Avengers?"
documents = SimpleDirectoryReader("./data").load_data()
node\_parser = VectorStoreIndex.from\_documents(documents)
query\_engine = node\_parser.as\_query\_engine()
response = query_engine.query(question)
print(f"base query result: {response}")
\# Output
base query result: No, Harley Quinn and Thanos are not depicted as righteous characters in the Avengers series.
上記のコードは、通常のRAG検索のためのLlamaIndexの使用方法を示しています。それをテストするために、ウィキペディアからアベンジャーズの映画の要約を使用しました。私たちはDCコミックスの「ハーレー・クイン」とマーベルの「サノス」についての複雑な質問をし、彼らがアベンジャーズで良いキャラクターであったかどうかを尋ねました。答えは正しかったが、一つのキャラクターが実はマーベルの映画からでは無いことを明確にしていませんでした。
さて、より小さな質問をどれだけうまく行うか見てみましょう：
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama\_index.core.query\_engine import SubQuestionQueryEngine
query\_engine\_tools = \[
QueryEngineTool(
query\_engine=query\_engine,
metadata=ToolMetadata(
name="Avengers",
description="Marvel movie The Avengers",
),
),
\]
query\_engine = SubQuestionQueryEngine.from\_defaults(
query\_engine\_tools=query\_engine\_tools
)
response = query_engine.query(question)
print(f"sub question query result: {response}")
\# Output
Generated 2 sub questions.
\[Avengers\] Q: What role does Harley Quinn play in the Avengers movie?
\[Avengers\] Q: What role does Thanos play in the Avengers movie?
\[Avengers\] A: Harley Quinn is not mentioned in the provided background of the Avengers movie.
\[Avengers\] A: Thanos is the primary antagonist in the Avengers movie. He is a powerful warlord who seeks to reshape the universe according to his vision. Thanos is depicted as a formidable and ruthless enemy, posing a significant threat to the Avengers and the entire universe.
sub question query result: Harley Quinn is not mentioned in the provided background of the Avengers movie. Thanos is the primary antagonist in the Avengers movie, depicted as a powerful and ruthless enemy.
- まず、検索ツールを構築し、その情報を設定します。
- 次に、検索ツールを使用して、大きな問題の小さな部分を処理できるツールをクラスで作成します。
- 結果は、小さな問題とその答えを示しています。最終的な答えは、これらすべての小さな答えに基づいています。
コードから、複雑な問題を分割すると結果が良くなることがわかります。
例では、デバッグのために小さなクエリとその答えが示されています。これらのデータも取得する際に取得できます：
from llama_index.core.callbacks import (
CallbackManager,
LlamaDebugHandler,
CBEventType,
EventPayload,
)
from llama_index.core import Settings
llama\_debug = LlamaDebugHandler(print\_trace\_on\_end=True)
callback\_manager = CallbackManager(\[llama\_debug\])
Settings.callback\_manager = callback\_manager
\# Subquestion querying code
...
for i, (start\_event, end\_event) in enumerate(
llama\_debug.get\_event\_pairs(CBEventType.SUB\_QUESTION)
):
qa\_pair = end\_event.payload\[EventPayload.SUB_QUESTION\]
print("Sub Question " \+ str(i) + ": " \+ qa\_pair.sub\_q.sub_question.strip())
print("Answer: " \+ qa_pair.answer.strip())
- サブクエスチョンのクエリに対するデバッグ情報を記録するために、コールバックマネージャーを追加します。
- クエリの後、コールバックマネージャーを通じてデバッグ情報を取得し、サブクエスチョンと回答を取得します。
プロンプトの生成
LlamaIndexは、サブクエスチョンを生成するための別のPythonパッケージ llama-index-question-gen-openaiを使用しています。内部的には、サブクエスションの生成にはデフォルトでOpenAIモデルを使用しています。プロンプトのテンプレートは、LlamaIndex公式リポジトリで見つけることができます。
LlamaIndexでプロンプトを表示するには、以下の方法があります。最初の方法は、get_prompts()メソッドを使用することです：
prompts = query\_engine.get\_prompts()
for key in prompts.keys():
sub\_question\_prompt = prompts\[key\]
template = sub\_question\_prompt.get_template()
print(f'prompt: {template}')
- まず、
get_prompts()メソッドを使用して、ほとんどのLlamaIndexオブジェクトで利用可能なクエリエンジンのプロンプトオブジェクトを取得します。 - プロンプトオブジェクトはJSONオブジェクトで、各キーがプロンプトテンプレートを表しています。
- プロンプトオブジェクトの各キーを順に処理し、各キーに対応するプロンプトテンプレートを取得し、それを出力します。
- サブクエスチョンのクエリは、サブクエスチョン生成用のプロンプトテンプレートと、通常のRAG用の別のプロンプトテンプレートを含みます。
もう一つの方法は、set_global_handlerを使用してグローバル設定を行うことです：
from llama_index.core import set\_global\_handler
set\_global\_handler("simple")
上記のコードをファイルの始めに追加すると、実行中にプロンプトが出力され、特定の変数値を持つ完全なプロンプトが表示されます。
HyDEクエリ変換
画像のソース: https://arxiv.org/pdf/2212.10496.pdf
HyDE（Hypothetical Document Embeddings）は、ユーザーのクエリに対して仮説的な文書を作成するためにLLMを使用します。これらの文書にはエラーが含まれているかもしれませんが、RAGの知識ベース内の文書を見つけるのに役立ちます。これらの仮説的な文書を実際のものと比較することで、より正確な結果を得ることができます。詳しくは、この論文をご覧ください。
コード例
HyDEを使用するLlamaIndexがどのように仮説的な文書を作成するかを見てみましょう:
from llama\_index.core.indices.query.query\_transform import HyDEQueryTransform
question = "What mysterious item did Loki use in an attempt to conquer Earth?"
hyde = HyDEQueryTransform(include_original=True)
query_bundle = hyde(question)
print(f"query_bundle embedding len: {len(query\_bundle.embedding\_strs)}")
for idx, embedding in enumerate(query\_bundle.embedding\_strs):
print(f"embedding {idx}: {embedding\[:100\]}")
\# Display result
query_bundle embedding len: 2
embedding 0: Loki used the Tesseract, also known as the Cosmic Cube, in his attempt to conquer Earth. This myste...
embedding 1: What mysterious item did Loki use in an attempt to conquer Earth?
- まず、
HyDEQueryTransformオブジェクトを作成します。オリジナルの質問を仮定のドキュメントに表示するために、include_original=Trueを設定します。include_originalはデフォルトでTrueですが、明確さのためにこのステップを表示しています。 - 次に、質問を使用して
hydeオブジェクトを呼び出します。これにより、QueryBundleオブジェクトが得られます。 -
QueryBundleオブジェクトにはembedding_strs属性があります。これは、仮定のドキュメントが最初の要素として配列になっています。include_originalがTrueの場合、オリジナルの質問は配列の2番目の要素になります。
LLMは、それが知っていることに基づいて質問に良い答えを提供できます。仮定のドキュメントは映画のプロットの要約と一致します。
以下は、LlamaIndexの仮定的なドキュメントのためのプロンプトテンプレートです。目標は、キーとなる詳細を含むパッセージで質問に答えることです。{context_str}はユーザーの質問を表します：
HYDE_TMPL = (
"Please write a passage to answer the question\\n"
"Try to include as many key details as possible.\\n"
"\\n"
"\\n"
"{context_str}\\n"
"\\n"
"\\n"
'Passage:"""\\n'
)
さあ、クエリエンジンを使って質問に答えを取得しましょう：
from llama\_index.core.query\_engine import TransformQueryEngine
hyde\_query\_engine = TransformQueryEngine(query_engine, hyde)
response = hyde\_query\_engine.query(question)
print(f"hyde query result: {response}")
\# Display result
hyde query result: Loki used the Tesseract in his attempt to conquer Earth. This powerful artifact, also kn...
- 初期のクエリエンジンと
HyDEQueryTransformを使用してTransformQueryEngineを作成します。 - エンジンの
query関数は、最初の質問から仮のドキュメントを作成し、その後、これらの仮のドキュメントから答えを見つけて作成します。
正しい答えが得られるものの、LlamaIndexが検索プロセス中にこれらの仮のドキュメントを使用しているのかどうかは不確かです。以下のコードでこれを確認できます：
from llama\_index.core.retrievers.transform\_retriever import TransformRetriever
retriever = node\_parser.as\_retriever(similarity\_top\_k=2)
hyde_retriever = TransformRetriever(retriever, hyde)
nodes = hyde_retriever.retrieve(question)
print(f"hyde retriever nodes len: {len(nodes)}")
for node in nodes:
print(f"node id: {node.id_}, score: {node.get_score()}")
print("=" \* 50)
nodes = retriever.retrieve("\\n".join(f"{n}" for n in query\_bundle.embedding\_strs))
print(f"hyde documents retrieve len: {len(nodes)}")
for node in nodes:
print(f"node id: {node.id_}, score: {node.get_score()}")
- 最初に、既存のリトリバーと
HyDEQueryTransformを使用して新しいリトリバーをTransformRetrieverで作成します。 - 次に、この新しいリトリバーを使用してユーザーの質問を見つけ、見つかったドキュメントのIDとスコアを表示します。
- 次に、最初のリトリバーを使用して推定されるドキュメントを探します。これらは
QueryBundleオブジェクトのembedding_strsから取得します。これには2つの部分があります: 一つは推定されるドキュメント、もう一つは元の質問です。 - これらの推定されるドキュメントのドキュメントIDとスコアを表示します。
以下が表示されます：
hyde retriever nodes len: 2
node id: 51e9381a-ef93-49ee-ae22-d169eba95549, score: 0.8895532276574978
node id: 5ef8a87e-1a72-4551-9801-ae7e792fdad2, score: 0.8499209871867581
==================================================
hyde documents retrieve nodes len: 2
node id: 51e9381a-ef93-49ee-ae22-d169eba95549, score: 0.8842142746289462
node id: 5ef8a87e-1a72-4551-9801-ae7e792fdad2, score: 0.8460828835028101
2つのアプローチの結果がほとんど同一であることがわかります。これは、検索に使用される入力が同じであることを示しています。すなわち、仮定のドキュメントです。HyDEQueryTransformオブジェクト内のinclude_original属性をFalseに設定して、生成された仮定のドキュメントが元の質問を含まないようにし、その後、コードを再度実行すると、結果は次のようになります：
hyde retriever nodes len: 2
node id: cfaea328-16d8-4eb8-87ca-8eeccad28263, score: 0.7548985780343257
node id: f47bc6c7-d8e1-421f-b9b8-a8006e768c04, score: 0.7508234876205329
==================================================
hyde documents retrieve nodes len: 2
node id: 6c2bb8cc-3c7d-4f92-b039-db925dd60d53, score: 0.7498683385309097
node id: f47bc6c7-d8e1-421f-b9b8-a8006e768c04, score: 0.7496147322045141
2つのやり方の結果は似ていますが、元の質問が足りないため、得られたドキュメントのスコアは低いです。
HyDEの制限
HyDEは、LLMの知識に基づいて仮説的なドキュメントを作成します。しかし、これらにはエラーや不正確さが含まれることがあります。公式ドキュメンテーションによると、HyDEはクエリを誤導し、バイアスをもたらすことがあります。ですので、実生活のタスクに使用する際には注意が必要です。
プロンプトステップバック
画像のソース：https://arxiv.org/pdf/2310.06117.pdf
ステップバックプロンプティングは、特定の例から複雑なアイデアを学習するための簡単な方法で、言語モデル（LLM）が論理的に思考し、問題を正確に解決するのに役立ちます。
例えば、上記の画像の最初の質問を見てみましょう。質問は温度と体積が与えられたときの圧力を尋ねています。最初の反応、直接的なものも思考過程のものも、間違っていました。しかし、ステップバックプロンプティングを使うと、まず元の質問から一般的な質問を作ります、例えば問題の背後にある物理学の公式を見つけるといったことです。この一般的な質問を解決し、この解答と元の質問をLLMに与えます。こうすることで、正しい答えを得ることができます。詳細は、ステップバックプロンプティングの論文こちらをご覧ください。
コード例
ステップバックのプロンプトはLlamaIndexに特別に実装されていませんが、オリジナルの呼び出しを通じてLLMとLlamaIndexを組み合わせることでそれを示すことができます。まず、LLMがオリジナルの質問に基づいたステップバックの質問を生成するようにしましょう。
from llama_index.core import PromptTemplate
from openai import OpenAI
client = OpenAI()
examples = \[
{
"input": "Who was the spouse of Anna Karina from 1968 to 1974?",
"output": "Who were the spouses of Anna Karina?",
},
{
"input": "Estella Leopold went to whichschool between Aug 1954and Nov 1954?",
"output": "What was Estella Leopold'seducation history?",
},
\]
few\_shot\_examples = "\\n\\n".join(
\[f"human: {example\['input'\]}\\nAI: {example\['output'\]}" for example in examples\]
)
step\_back\_question\_system\_prompt = PromptTemplate(
"You are an expert at world knowledge."
"Your task is to step back and paraphrase a question to a more generic step-back question,"
"which is easier to answer. Here are a few examples:\\n"
"{few\_shot\_examples}"
)
completion = client.chat.completions.create(
model="gpt-3.5-turbo",
temperature=0.1,
messages=\[
{
"role": "system",
"content": step\_back\_question\_system\_prompt.format(
few\_shot\_examples=few\_shot\_examples
),
},
{"role": "user", "content": question},
\],
)
step\_back\_question = completion.choices\[0\].message.content
print(f"step\_back\_question: {step\_back\_question}")
- まず、ステップバックの質問の例をいくつか定義し、それらをLLMのシステムプロンプトに含めて、質問生成のパターンを理解させます。
- 次に、ユーザーの質問とシステムプロンプトの両方をLLMに提供して、ステップバックの質問を生成します。
ステップバックの質問を生成した後、元の質問とステップバックの質問の両方に関連するドキュメントを別々に取得します:
retrievals = retriever.retrieve(question)
normal_context = "\\n\\n".join(\[f"{n.text}" for n in retrievals\])
retrievals = retriever.retrieve(step\_back\_question)
step\_back\_context = "\\n\\n".join(\[f"{n.text}" for n in retrievals\])
検索結果を得たら、LLMに最終的な答えを生成させます：
step\_back\_qa\_prompt\_template = PromptTemplate(
"Context information is below.\\n"
"---------------------\\n"
"{normal_context}\\n"
"{step\_back\_context}\\n"
"---------------------\\n"
"Given the context information and not prior knowledge, "
"answer the question: {question}\\n"
)
completion = client.chat.completions.create(
model="gpt-3.5-turbo",
temperature=0.1,
messages=\[
{
"role": "system",
"content": "Always answer the question, even if the context isn't helpful.",
},
{
"role": "user",
"content": step\_back\_qa\_prompt\_template.format(
normal\_context=normal\_context,
step\_back\_context=step\_back\_context,
question=question,
),
},
\],
)
step\_back\_result = completion.choices\[0\].message.content
print(f"step\_back\_result: {step\_back\_result}")
- プロンプトテンプレートでは、元の質問とステップバックの質問の両方に関連する文書情報をLLMに提供し、それらを元の質問と組み合わせてLLMに答えを生成させます。
最後に、ステップバックプロンプティングの有無による通常のRAG検索の結果を比較しましょう：
question: Was there a great war on the planet Titan??
base result: No, there has not been a major war on the planet Titan. It is not known for being the site of any significant conflicts or wars.
====================================================================================================
step back question: Have there been any significant events on the planet Titan?
step back result: Yes, in the Marvel Cinematic Universe, there was a significant conflict on the planet Titan. In "Avengers: Infinity War," Titan is depicted as the destroyed homeworld of Thanos, and the battle on Titan involved a group of heroes including Iron Man (Tony Stark), Spider-Man (Peter Parker), Doctor Strange (Stephen Strange), and the Guardians of the Galaxy, as they attempted to thwart Thanos from achieving his goals.
ステップバックプロンプティングを使用せずに結果が間違っていることがわかります。しかし、ステップバックプロンプティングを使用すると、知識に基づいて正しい答えを得ることができます。
まとめ
今日は、RAG検索でクエリを書き換えるためのいくつかの異なる戦略について話しました。LlamaIndexのコード例を用いてこれらの戦略を示し、LlamaIndexの使用に関するいくつかのヒントを提供しました。本記事では触れなかった他の戦略もあります。RAG技術が進化すると、戦略も増えてきます。その時が来たら、この部分にもっと追加します。