编程开源技术交流,分享技术与知识

网站首页 > 开源技术 正文

高级 RAG 检索策略:递归检索(高级检索方法的特点)

wxchong 2024-08-16 06:09:18 开源技术 11 ℃ 0 评论

随着 LLM(大型语言模型)技术的进步,RAG(检索增强生成)技术在问答和对话等任务中变得越来越普遍。RAG 技术的一个关键组成部分是文档检索器,它负责从庞大的数据集中检索与查询相关的文档,以帮助 LLM 生成答案。RAG检索器的有效性直接影响LLM的响应质量,使得高效RAG检索器的设计成为一个重要的研究课题。目前,RAG检索有多种策略。本文将介绍一种先进的 RAG 检索策略——递归检索,它通过递归过程检索相关文档,从而提高检索效率。

递归检索简介

与标准 RAG 检索相比,递归检索解决了后者文档段过大导致的不准确性问题。下面是递归检索过程的流程图:

  • 递归检索扩展了具有较小粒度节点的原始文档节点
  • 在文档检索过程中,如果检索到扩展的节点,则进程会递归检索其原始节点,然后将原始节点作为检索结果提交给 LLM

在 LlamaIndex 的实现中,递归检索主要以两种形式发生:块引用递归检索和元数据引用递归检索。

标准 RAG 检索

在深入研究递归检索之前,让我们看一个使用 LlamaIndex 的标准 RAG 检索示例:

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import VectorStoreIndex

question = "Which two members of the Avengers created Ultron?"
documents = SimpleDirectoryReader("./data").load_data()
node_parser = SentenceSplitter(chunk_size=1024)
base_nodes = node_parser.get_nodes_from_documents(documents)
print(f"base_nodes len: {len(base_nodes)}")
for idx, node in enumerate(base_nodes):
    node.id_ = f"node-{idx}"
base_index = VectorStoreIndex(nodes=base_nodes)
base_retriever = base_index.as_retriever(similarity_top_k=2)
retrievals = base_retriever.retrieve(question)
for n in retrievals:
    print(
        f"Node ID: {n.node_id}\nSimilarity: {n.score}\nText: {n.text[:100]}...\n"
    )
response = base_retriever.query(question)
print(f"response: {response}")
print(f"len: {len(response.source_nodes)}")
  • 我们将维基百科上复仇者联盟电影情节中的文档放在文档测试数据的数据目录中
  • SentenceSplitter 文档解析器分析文档,尽可能保持句子和段落的完整性,默认chunk_size为 1024
  • 解析后,原始节点 ID 默认为随机字符串,我们将其格式化为 node-{idx},以便以后验证检索结果
  • 然后,我们创建一个 VectorStoreIndex 索引,输入原始节点,并创建一个 similarity_top_k=2 的检索器base_retriever,这意味着它在检索过程中按相似性返回前 2 个节点,然后打印出检索到的节点信息
  • 最后,检索器生成问题的答案并打印出来

让我们看看该计划的结果:

base_nodes len: 15
Node ID: node-0
Similarity: 0.8425314373498192
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...

Node ID: node-1
Similarity: 0.8135015554872678
Text: In 2018, twenty-three days after Thanos erased half of all life in the universe,[a] Carol Danvers re...

response: Tony Stark and Bruce Banner
nodes len: 2

我们可以看到,经过文档解析器解析后,有 15 个原始节点,其中 2 个节点被检索到,都是原始节点。

块引用递归检索

块引用递归检索建立在标准 RAG 检索的基础上,将每个原始文档节点拆分为较小的文档节点,这些节点与原始节点保持父子关系。当检索到子节点时,它会递归地检索其父节点,然后将父节点作为检索结果提交给 LLM。

下面,我们使用一个代码示例来理解块引用递归检索,首先创建几个具有较小chunk_size值的文档解析器:

sub_chunk_sizes = [128, 256, 512]
sub_node_parsers = [
    SentenceSplitter(chunk_size=c, chunk_overlap=20) for c in sub_chunk_sizes
]

然后,我们使用文档解析器将原始节点解析为子节点:

from llama_index.core.schema import IndexNode

all_nodes = []
for base_node in base_nodes:
    for n in sub_node_parsers:
        sub_nodes = n.get_nodes_from_documents([base_node])
        sub_inodes = [
            IndexNode.from_text_node(sn, base_node.node_id) for sn in sub_nodes
        ]
        all_nodes.extend(sub_inodes)
    original_node = IndexNode.from_text_node(base_node, base_node.node_id)
    all_nodes.append(original_node)
print(f"all_nodes len: {len(all_nodes)}")
# Display result
all_nodes len: 331
  • 我们使用每个小块文档解析器来解析原始节点,然后将解析后的子节点和原始节点放入all_nodes列表中
  • 每个原始节点的chunk_size为 1024,当除以 512 的chunk_size时,生成大约 2 个子节点;256 chunk_size生成大约 4 个子节点;chunk_size 128 生成大约 8 个子节点
  • 每个子节点的 node_id 属性值是原始节点的id_,我们之前将其格式化为 node-{idx},但子节点的 id_ 属性值仍然是 LlamaIndex 生成的随机字符串
  • 原始节点是 TextNode 类型的节点,我们将其转换为 IndexNode 类型的节点并添加到all_nodes列表中,最终生成 331 个节点

接下来,我们创建一个包含所有节点的检索索引,执行标准检索,并观察结果:

vector_index_chunk = VectorStoreIndex(all_nodes)
vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=2)
nodes = vector_retriever_chunk.retrieve(question)
for node in nodes:
    print(
        f"Node ID: {node.node_id}\nSimilarity: {node.score}\nText: {node.text[:100]}...\n"
    )

# Display result
Node ID: 0e3409e5-6c84-4bbf-886a-40e8553eb463
Similarity: 0.8476561735049716
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...

Node ID: 0ed2ca24-f262-40fe-855b-0eb84c1a1567
Similarity: 0.8435371049710689
Text: They meet two of Strucker's test subjects—twins Pietro (who has superhuman speed) and Wanda Maximoff...
  • 我们创建一个 VectorStoreIndex 索引,输入所有节点,并创建一个 similarity_top_k=2 的检索器vector_retriever_chunk,表示检索按相似度返回前 2 个节点
  • 在标准检索结果中,我们可以看到检索了两个子节点,因为它们的节点 ID 是一个随机字符串,而不是我们之前格式化的 node-{idx}

接下来,让我们使用递归检索来检查结果:

from llama_index.core.retrievers import RecursiveRetriever

all_nodes_dict = {n.node_id: n for n in all_nodes}
retriever_chunk = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_chunk},
    node_dict=all_nodes_dict,
    verbose=True,
)
nodes = retriever_chunk.retrieve(question)
for node in nodes:
    print(
        f"Node ID: {node.node_id}\nSimilarity: {node.score}\nText: {node.text[:1000]}...\n"
    )

# Display result
Retrieving with query id None: Which two members of the Avengers created Ultron?
Retrieved node with id, entering: node-0
Retrieving with query id node-0: Which two members of the Avengers created Ultron?
Node ID: node-0
Similarity: 0.8476561735049716
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...
  • 我们首先构造一个all_nodes_dict字典,将所有节点的node_id作为键,以节点对象为值,从而实现node_id递归检索
  • 然后,我们创建一个 RecursiveRetriever 检索器,输入 vector_retriever_chunk 检索器并all_nodes_dict字典,verbose=True 打印检索过程
  • 最后,我们对问题进行递归检索,我们可以看到检索结果是单个原始节点,因为之前的标准检索结果返回了两个子节点,其父节点是相同的原始节点,因此在递归检索期间只返回该原始节点,并且该节点的相似性得分与标准检索结果中的第一个节点匹配: 0.8476561735049716

最后,我们使用 LLM 生成一个答案:

from llama_index.core.query_engine import RetrieverQueryEngine

llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
query_engine_chunk = RetrieverQueryEngine.from_args(retriever_chunk, llm=llm)
response = query_engine_chunk.query(question)
print(f"response: {str(response)}")
print(f"nodes len: {len(response.source_nodes)}")

# Display result
response: Tony Stark and Bruce Banner
nodes len: 1

我们可以看到递归检索生成的答案与标准 RAG 检索生成的答案相同。

元数据引用递归检索

元数据引用递归检索类似于块引用递归检索,只不过它不拆分原始节点。相反,它基于原始节点生成元数据子节点,然后将元数据子节点和原始节点都输入到检索索引中。

下面,我们使用一个代码示例来了解元数据引用递归检索,首先创建几个元数据提取器:

from llama_index.core.extractors import (
    SummaryExtractor,
    QuestionsAnsweredExtractor,
)

extractors = [
    SummaryExtractor(summaries=["self"], show_progress=True),
    QuestionsAnsweredExtractor(questions=5, show_progress=True),
]
  • 我们创建了两个元数据提取器,一个是用于生成文档摘要的 SummaryExtractor,另一个是用于生成文档可以回答的问题的 QuestionsAnsweredExtractor
  • QuestionsAnsweredExtractor 的 questions=5 参数表示它生成 5 个问题
  • show_progress=True 表示提取过程
  • 这两个提取器使用 LLM 生成元数据,默认为 OpenAI 的 GPT-3.5-turbo 模型

然后,我们使用元数据提取器将原始节点解析为元数据子节点:

node_to_metadata = {}
for extractor in extractors:
    metadata_dicts = extractor.extract(base_nodes)
    for node, metadata in zip(base_nodes, metadata_dicts):
        if node.node_id not in node_to_metadata:
            node_to_metadata[node.node_id] = metadata
        else:
            node_to_metadata[node.node_id].update(metadata)
  • 我们使用两个提取器为原始节点生成元数据,将结果保存在node_to_metadata字典中
  • node_to_metadata 的键是原始文档的node_id,其值是原始节点的元数据,包括摘要和问题

代码执行后node_to_metadata的数据结构如下图所示:

{
  "node-0": {
    "section_summary": "...",
    "questions_this_excerpt_can_answer": "1. ...?\n2. ...?\n3. ...?\n4. ...?\n5. ...?"
  },
  "node-1": {
    "section_summary": "...",
    "questions_this_excerpt_can_answer": "1. ...?\n2. ...?\n3. ...?\n4. ...?\n5. ...?"
  },
  ......
}

我们可以将node_to_metadata数据保存到文件中以备后用,因此我们不需要每次都调用 LLM 来生成元数据。

import json

def save_metadata_dicts(path, data):
    with open(path, "w") as fp:
        json.dump(data, fp)

def load_metadata_dicts(path):
    with open(path, "r") as fp:
        data = json.load(fp)
    return data
save_metadata_dicts("output/avengers_metadata_dicts.json", node_to_metadata)
node_to_metadata = load_metadata_dicts("output/avengers_metadata_dicts.json")
  • 我们定义了两种方法,save_metadata_dicts用于将元数据字典保存到文件,load_metadata_dicts用于从文件加载元数据字典
  • 我们将元数据字典保存到 avengers_metadata_dicts.json 文件中
  • 稍后,当我们需要再次使用元数据字典时,我们可以使用 load_metadata_dicts 方法直接从文件中加载它

然后,我们将原始节点和元数据子节点合并到一个新的节点列表中:

import copy

all_nodes = copy.deepcopy(base_nodes)
for node_id, metadata in node_to_metadata.items():
    for val in metadata.values():
        all_nodes.append(IndexNode(text=val, index_id=node_id))
print(f"all_nodes len: {len(all_nodes)}")

# Display result
all_nodes len: 45
  • 我们首先将原始节点复制到新的节点列表中
  • 然后,我们将元数据字典的摘要和问题作为新节点添加到新节点列表中,将它们与原始节点相关联以形成父子关系
  • 此过程生成 45 个节点,包括 15 个原始节点和 30 个元数据子节点

我们可以在新节点列表中检查 node-0 原始节点及其子节点的内容:

node0_nodes = list(
    filter(
        lambda x: x.id_ == "node-0"
        or (hasattr(x, "index_id") and x.index_id == "node-0"),
        all_nodes,
    )
)
print(f"node0_nodes len: {len(node0_nodes)}")
for node in node0_nodes:
    index_id_str = node.index_id if hasattr(node, 'index_id') else 'N/A'
    print(
        f"Node ID: {node.node_id}\nIndex ID: {index_id_str}\nText: {node.text[:100]}...\n"
    )

# Display result
node0_nodes len: 3
Node ID: node-0
Index ID: N/A
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...

Node ID: 45d41128-a8e6-4cdc-8ef3-7a71f01ddd96
Index ID: node-0
Text: The key topics of the section include the creation of Ultron by Tony Stark and Bruce Banner, the int...

Node ID: a06f3bb9-8a57-455f-b0c6-c9602b107158
Index ID: node-0
Text: 1. What are the names of the Avengers who raid the Hydra facility in Sokovia at the beginning of the...
  • 我们使用 filter 函数过滤掉 node-0 的原始节点及其关联的元数据子节点,共计 3 个节点
  • 第一个是原始节点,第二个是元数据摘要子节点,第三个是元数据问题子节点

接下来,我们创建一个检索索引,输入所有节点,执行标准检索,并观察结果:

vector_index_metadata = VectorStoreIndex(all_nodes)
vector_retriever_metadata = vector_index_metadata.as_retriever(similarity_top_k=2)

enginer = vector_index_metadata.as_query_engine(similarity_top_k=2)
nodes = enginer.retrieve(question)
for node in nodes:
    print(
        f"Node ID: {node.node_id}\nSimilarity: {node.score}\nText: {node.text[:100]}...\n"
    )

# Display result
Node ID: d2cc032a-b258-4715-b335-ebd1cf80494d
Similarity: 0.857976008616706
Text: The key topics of the section include the creation of Ultron by Tony Stark and Bruce Banner, the int...

Node ID: node-0
Similarity: 0.8425314373498192
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...
  • 我们创建一个 VectorStoreIndex 索引,输入所有节点,并创建一个 similarity_top_k=2 的检索器vector_retriever_metadata,表示检索按相似度返回前 2 个节点
  • 在标准检索结果中,我们可以看到检索了两个节点,第一个是元数据摘要子节点,第二个是原始节点,可通过其节点 ID 识别

以上是标准检索结果。接下来,让我们使用递归检索来检查结果:

all_nodes_dict = {n.node_id: n for n in all_nodes}
retriever_metadata = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_metadata},
    node_dict=all_nodes_dict,
    verbose=False,
)
nodes = retriever_metadata.retrieve(question)
for node in nodes:
    print(
        f"Node ID: {node.node_id}\nSimilarity: {node.score}\nText: {node.text[:100]}...\n\n"
    )

# Display result
Node ID: node-0
Similarity: 0.857976008616706
Text: In the Eastern European country of Sokovia, the Avengers—Tony Stark, Thor, Bruce Banner, Steve Roger...
  • 这里的代码类似于前面的块引用递归检索,只是用vector_retriever_metadata替换vector_retriever_chunk,然后对问题执行递归检索
  • 我们可以看到最终的检索结果是一个原始节点,因为之前的标准检索结果返回了一个元数据子节点和一个原始节点,而这个子节点的父节点是同一个原始节点,因此在递归检索期间只返回这个原始节点,这个原始节点的相似性分数与标准检索结果中的第一个节点匹配: 0.857976008616706

最后,我们使用 LLM 生成一个答案:

query_engine_metadata = RetrieverQueryEngine.from_args(retriever_metadata, llm=llm)
response = query_engine_metadata.query(question)
print(f"response: {str(response)}")
print(f"nodes len: {len(response.source_nodes)}")

# Display result
response: Tony Stark and Bruce Banner
nodes len: 1

我们可以看到递归检索生成的答案与标准 RAG 检索生成的答案相同。

检索效果比较

接下来,我们使用 Trulens 来评估标准 RAG 检索、块引用递归检索和元数据引用递归检索的有效性。

tru.reset_database()
rag_evaluate(base_engine, "base_evaluation")
rag_evaluate(engine, "recursive_retriever_chunk_evaluation")
rag_evaluate(engine, "recursive_retriever_metadata_evaluation")
Tru().run_dashboard()

rag_evaluate的具体代码可以在我之前的文章中找到,主要使用 Trulens 的接地性、qa_relevance和qs_relevance来评估 RAG 检索结果。执行完代码后,我们可以在浏览器中查看 Trulens 中的评估结果:

在评估结果中,我们可以看到两种类型的递归检索都比标准 RAG 检索性能更好,元数据引用递归检索的性能略优于块引用递归检索。但是,评估结果不是绝对的,实际效果需要根据具体情况进行评估。

总结

递归检索是一种高级 RAG 检索策略,它首先将原始文档节点扩展为更精细的文档节点,从而在检索过程中更准确地检索相关文档,然后使用递归检索查找匹配的原始文档节点。递归检索可以增强RAG检索的有效性,但也会增加检索时间和计算资源,因此应根据具体情况选择合适的检索策略。

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表