作者： 极客与黑客之路 来源： 极客与黑客之路

之前介绍了高级 RAG 检索的句子窗口检索策略，今天我们再来介绍另外一种高级检索策略——自动合并检索，它比句子窗口要复杂一些，但请不用担心，下面的介绍会让你理解其中原理，同时会介绍如何使用 LlamaIndex 来构建一个自动合并检索，最后使用 Trulens 来对检索效果进行评估，并与之前的检索策略进行对比。

自动合并检索介绍

自动合并检索主要是将文档按照块大小拆分成不同层级的节点，这些节点包括父节点和子节点，然后在检索过程中找到相似度高的叶子节点，如果一个父节点中有多个子节点被检索到，那么这个父节点就会被自动合并，最终将父节点的所有文档都作为上下文发送给 LLM（大语言模型），下面是自动合并检索的示意图：

自动合并检索是 LlamaIndex 中的一种高级检索功能，主要有文档拆分和文档合并两个过程，下面我们将通过代码来讲解其中的原理。

文档拆分

在构建一个自动合并检索时，我们首先要创建一个 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，第二个层级的文档大小为 512，第三个层级的文档大小为 128。当然你也可以将层级设置为更少或者更多，比如设置成 2 级，那么chunk_sizes 可以是[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)}")  
  
# 显示结果  
total len: 66  
root len: 4  
leaf len: 52  

• get_leaf_nodes 和get_root_nodes 这 2 个方法都是传入一个节点列表

• 可以看到总的节点数是 66，根节点是 4，叶子节点是 52

• 根节点加上叶子节点的总数是 56（4+52），和总节点数 66 并不匹配，所以剩下的节点是中间层级的节点，我们可以推算出中间节点数是 10（66-56）

• 如果你的文档层级是 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)}")  
  
# 显示结果  
deep0 len: 4  
deep1 len: 10  
deep2 len: 52  

• get_deeper_nodes 方法第一个参数是节点列表，第二参数是要查询的层级，0 表示第 1 层级，也就是根节点

可以看到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)}")  
  
# 显示结果  
middle len: 10  
leaf len: 52  

• get_child_nodes 方法第一个参数是要获取子节点的节点列表，第二个参数是所有节点

• 这里我们先获取根节点下的所有子节点，得到 10 个子节点，这些节点也就是中间层级节点

• 然后我们再获取这些中间节点下的所有子节点，得到 52 个子节点，这些节点也就是叶子节点

当然我们也可以获取某个节点下的子节点，比如获取第一个根节点的子节点：

root0_child_nodes = get_child_nodes(root_nodes[0], all_nodes=nodes)  
print(f"root0 child len: {len(root0_child_nodes)}")  
  
# 显示结果  
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}")  
  
# 显示结果  
deep1[0] node: 自从宇宙魔方于1942年被人类发现后针对其展开过无数次探索，栖息于宇宙中的强大外星势力也从此开始盯住地球，被它们营救的洛基负责率领齐塔瑞军团。在地球，由神盾局建立的一所科研基地[注 14]中进行着神盾局与美国国家航空航天局和美国空军合作的“天马项目”[注 15]，试图提炼出魔方的能量并加以利用，但却造成魔方的能量数值持续攀升。神盾局探员菲尔·考森和玛丽亚·希尔受尼克·弗瑞局长的命令疏散基地，而魔方却于地下实验室自行开启传送门将洛基传送过来。洛基杀光所有护卫后，用他手上的一把能灌输能量的权杖洗脑并操纵弗瑞的亲信克林特·巴顿和协助神盾局着手魔方项目的科学家埃里克·塞尔维格格，在他们的陪同下带着魔方坐车逃离基地。没过多久，流出的能量爆发造成基地完全坍塌，弗瑞为了应对外来威胁而命令所有人备战。  
child[0] node of deep1[0]: 自从宇宙魔方于1942年被人类发现后针对其展开过无数次探索，栖息于宇宙中的强大外星势力也从此开始盯住地球，被它们营救的洛基负责率领齐塔  

• 我们首先打印中间层级第一个节点的文档内容

• 然后再获取这个中间节点第一个子节点，并打印其文档内容

• 可以看到父节点的文档内容包含了子节点的文档内容

文档合并

文档合并是自动合并检索的重要组成部分，文档合并的效果决定了提交给 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 项目中的使用，文档数据我们还是使用之前维基百科上的复仇者联盟[1]电影剧情来进行测试。

自动合并检索示例

我们来看下如何使用 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 和 Embedding 模型进行答案生成和向量生成

• 再创建storage_context 来保存所有节点nodes ，后面的自动合并检索会根据叶子节点来找其相关的父节点，所以这里需要保存所有节点

• 接下来我们先构建一个基础检索base_index ，这个检索会根据问题对所有叶子节点leaf_nodes 进行检索，找到匹配度最高的similarity_top_k 个节点，这里我们将获取 12 个匹配度最高的叶子节点

• 我们再构建一个自动合并检索AutoMergingRetriever ，这个检索会根据基础检索的结果来进行合并操作，这里我们设置了simple_ratio_thresh 为 0.3，即当检索子节点比例大于这个阀值的节点就会进行自动合并。verbose 参数设置为 True，表示输出合并的过程

• 最后我们使用RetrieverQueryEngine 来创建一个检索引擎

接下来我们就可以使用这个检索引擎来回答问题了：

question = "奥创是由哪两位复仇者联盟成员创造的？"  
response = auto_merging_engine.query(question)  
print(f"response: {str(response)}")  
print(f"nodes len: {len(response.source_nodes)}")  
  
# 显示结果  
> Merging 5 nodes into parent node.  
> Parent node id: 80d1eeed-3447-4987-b05d-49fd4b6aabd4.  
> Parent node text: 神盾局解散后，由托尼·斯塔克、史蒂芬·罗杰斯、雷神、娜塔莎·罗曼诺夫、布鲁斯·班纳以及克林特·巴顿组成的复仇者联盟负责全力搜查九头蛇的下落，这次透过“盟友”提供的情报而进攻位于东欧的国家“索科维...  
  
> Merging 4 nodes into parent node.  
> Parent node id: 2e719ad1-89fe-4d00-add4-e0296b19eab3.  
> Parent node text: 复仇者们到达后跟他们正面交锋，但大多数人被旺达用幻象术迷惑，看到各自心中最深层的“阴影”；唯独托尔看见在家乡阿萨神域发生的不明景象。旺达同时迷惑班纳的大脑，使其丧失理智而变成绿巨人跑到约翰内斯堡...  
  
> Merging 2 nodes into parent node.  
> Parent node id: c1e7e8a1-d50b-4a35-9b0d-beec29993d1a.  
> Parent node text: 奥创发布机械大军，在位于城市正中央的教堂里启动靠振金造的陆地合成器，使整座城市陆地上升，企图透过陨石撞击方式灭绝人类后由机器人取代。班纳潜入基地救出娜塔莎，娜塔莎将他从高处推落释放出绿巨人支持战...  
  
response: 奥创是由托尼·斯塔克和布鲁斯·班纳这两位复仇者联盟成员创造的。  
nodes len: 4  

在没有经过自动合并之前，我们让基础检索获取了 12 个匹配度最高的叶子节点，在输出结果中可以看到，这 12 个节点经过了 3 次合并操作，最终我们得到了 4 个节点，这些节点中既包含叶子节点，也包含合并过后的父节点。

检索效果对比

我们再使用Trulens[2]来评估自动合并检索的效果：

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 的评估结果：

在评估结果中，我们可以看到自动合并检索相比其他两种检索的效果要好，但这不表示自动合并检索会一直比其他检索好，具体的评估效果还要看原始的输入文档，以及检索的参数设置等，总之，具体的评估效果要根据实际情况来评估。

总结

自动合并检索是高级 RAG 检索的一种方法，文档拆分和文档合并的思想是该方法的主要特点，本文介绍了自动合并检索的原理和实现方法，并使用 Trulens 来评估了自动合并检索的效果，希望可以帮助大家更好地理解和使用自动合并检索。

关注我，一起学习各种人工智能和 AIGC 新技术，欢迎交流，如果你有什么想问想说的，欢迎在评论区留言。

参考:

[1] 复仇者联盟: https://en.wikipedia.org/wiki/Avenger

[2] Trulens: https://www.trulens.org/

更多AI工具，参考Github-AiBard123，国内AiBard123