高级RAG(八): 自动合并检索(Auto-merging Retrieval)

自动合并检索(Auto-merging Retrieval)是LlamaIndex的另外一种高级RAG技术，它有点类似与我们之间介绍的从小到大的检索，不过自动合并检索要比“从小到大的检索”稍微复杂一些，它首先将文档按一定的层次结构进行切割，然后在检索的时候将被切割的最小结构也就是叶子节点与问题进行相似度匹配，当一个父节点上的多数叶子节点都与问题匹配上，则将该父节点文档作为context返回给LLM，下面我们先来介绍一下文档的层次结构：

一、文档的层次结构

所谓的文档层次结构是指当我们在切割文档的时候可以按特定的层次结构来进行切割，这里的层次结构可以理解为广度和深度两个维度，我们首先将一个文档切割成若干个较大的文档块，那么这些较大的文档块之间可以理解为是“兄弟关系”，然后我们将每个大文档块再切割成若干个较小的文档块，那么大文档块和小文档块之间就是“父子关系”，而这里的“兄弟关系”可以理解为广度，“父子关系”则可以理解为深度，为了进一步加深印象，我们可以这样理解：一个长文档可以由若干个章节组成的；每个章节又由多个段落组成，每个段落又可以进一步分割成多个句子，那么章节之间，段落之间，句子之间就是兄弟关系;章节和段落之间，段落和句子之间就是父子关系。下面我们借用github上的DocArray项目来说明文档的层次结构如下图所示：

二、自动合并检索

LlamaIndex的自动合并检索时首先需要将文档按特定的层次结构进行切割，因此先要定义切割的层次结构如[1024,512], 这表示将文档按两层结构进行切割即首先将文档按块大小(chunk_size)为1024进行切割，切割成若干个大文档块，然后每个大文档块(chunk_size=1024)再被切分成4个块大小为512的子文档块，那么这些子文档块就是所谓的叶子节点，而子文档块所属的大文档块就是所谓的父节点，再比如定义切割的层次结构为[1024,512,128],这表示将文档按三层结构进行切割，那么顶层节点的块大小为1024，中间层的块大小为512，底层的叶子节点的块大小为128。而在检索时只拿叶子节点和问题进行匹配，当某个父节点下的多数叶子节点都与问题匹配上则将该父节点作为context返回给LLM,如下图所示：

接下来就是我们的实战环节，我们仍然使用上一篇博客中的使用的百度百科的文档，这次我们会使用LlamaIndex的AutoMergingRetriever对文档进行检索，最后我们仍然使用Trulens对检索结果进行评估。下面我们首先进行环境配置。

一、环境配置

我们首先需要安装如下python包

pip install llama_hub 
pip install llama_index
pip install trulens-eval
pip install trafilatura
pip install torch sentence-transformers

接下来我们需要做一些初始化的工作，比如导入openai的api_key:

import os
 
os.environ['OPENAI_API_KEY']="your_api_key"

二、加载数据

这里我们仍然使用前几篇博客中使用的数据即从百度百科的网页中抓取一篇关于恐龙的科普文章:

下面我们使用LlamaIndex的网页爬虫工具TrafilaturaWebReader来爬取百度百科上的这篇关于恐龙的科普文章：

from llama_index.readers.web import TrafilaturaWebReader

url="https://baike.baidu.com/item/恐龙/139019"
documents = TrafilaturaWebReader().load_data([url])
documents

三、设置文档层次结构

在我们进行自动合并检索之前，我们需要创建一个文档切割器，它会按指定的文档的层次结构来切割文档，这里我们会将文档的层次结构设置为3层，文档切割器在切割文档时就会采样递归的切割方式对文档进行切割：

from llama_index.node_parser import HierarchicalNodeParser

node_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]
)

这里我们设置了文档的层次结构为[2048, 512, 128]，这就意味着每个叶子节点的大小(chunk_size)为128，需要说明的是这里的128是只文档的token数为128。下面我们就来使用这个文档切割器来切割之前下载的关于恐龙的科普文章：

nodes = node_parser.get_nodes_from_documents(documents)

len(nodes)

这里我们看到原来的文档被切割成了337个文档块，接下来我们来看一下其中的某一个文档块的内容：

nodes[50]

这里我们需要说明的是我们查看的是第50个文档块的内容，其中的NodeRelationship属性包含有SOURCE，PREVIOUS，NEXT，PARENT，CHILD，text 等与其相关的节点的索引信息(node_id),下面我们来解释一下这些相关节点属性的含义：

SOURCE: 指当前文档块源自于哪个文档(同PARENT)
PREVIOUS：指当前文档块的前一篇文档(与当前文档是兄弟关系)。
NEXT：指当前文档块的后一篇文档(与当前文档是兄弟关系)。
PARENT：指当前文档块所属的父文档(与当前文档是父子关系,同SOURCE)。
CHILD：指当前文档块的子文档(与当前文档是父子关系)。
text：当前文档块的内容。

下面我们可以单独查看叶子节点文档：

from llama_index.node_parser import get_leaf_nodes

leaf_nodes = get_leaf_nodes(nodes)
leaf_nodes[30]

这里我们查看了第30个叶子节点文档，由于是叶子节点，所以它不存在CHILD信息。同样也可以查看该叶子节点文档所属的父文档信息：

nodes_by_id = {node.node_id: node for node in nodes}

parent_node = nodes_by_id[leaf_nodes[30].parent_node.node_id]
parent_node

这里我们看到第30个叶子节点的父文档中的text信息中包含了子节点的所有文本内容。

四、创建向量库索引

文档切割完成以后，接下来我们需要创建向量库索引，其步骤类似与我之前写的博客：句子-窗口检索，不熟悉的朋友可以看一下这篇博客。

from llama_index.llms import OpenAI
from llama_index import ServiceContext
from llama_index import VectorStoreIndex, StorageContext

#创建LLM
llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)

#创建ServiceContext
auto_merging_context = ServiceContext.from_defaults(
    llm=llm,
    embed_model="local:BAAI/bge-small-zh-v1.5",
    node_parser=node_parser,
)

#创建向量库索引
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)
automerging_index = VectorStoreIndex(
    leaf_nodes, 
    storage_context=storage_context, 
    service_context=auto_merging_context
)

#向量库持久化
automerging_index.storage_context.persist(persist_dir="./merging_index")

这里我们使用的LLM是Openai的gpt-3.5-turbo，embedding模型为bge-small-zh-v1.5，最后我们做了向量库的持久化操作即将向量库保存在本地，这样做的好处是下次使用时不需要重新去下载数据和创建向量库索引，而只需读取本地保存的向量库索引就可以了。下面我们就来测试一下读取刚才保存在本地的向量库索引(这个步骤是可选的)：

import os
from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage
from llama_index import load_index_from_storage

if not os.path.exists("./merging_index"):
    storage_context = StorageContext.from_defaults()
    storage_context.docstore.add_documents(nodes)

    automerging_index = VectorStoreIndex(
            leaf_nodes,
            storage_context=storage_context,
            service_context=auto_merging_context
        )

    automerging_index.storage_context.persist(persist_dir="./merging_index")
else:
    automerging_index = load_index_from_storage(
        StorageContext.from_defaults(persist_dir="./merging_index"),
        service_context=auto_merging_context
    )

五、定义和执行检索器

创建了向量库索引以后，我们就需要创建检索器，检索引擎，然后通过该检索引擎对用户问题进行检索了：

from llama_index.indices.postprocessor import SentenceTransformerRerank
from llama_index.retrievers import AutoMergingRetriever
from llama_index.query_engine import RetrieverQueryEngine

base_retriever = automerging_index.as_retriever(
    similarity_top_k=12
)

retriever = AutoMergingRetriever(
    base_retriever, 
    automerging_index.storage_context, 
    verbose=True
)

rerank = SentenceTransformerRerank(top_n=6, model="BAAI/bge-reranker-base")

auto_merging_engine = RetrieverQueryEngine.from_args(
    retriever, node_postprocessors=[rerank]
)

这里我们创建了一个自动合并检索器automerging_retriever，它有一个输入参数similarity_top_k，我们将其设置为12，这意味着检索器每次在检索时会返回12个相关文档(context), 同时我们还设置了一个rerank模型bge-reranker-base以及该模型的参数top_n为6，rerank模型的作用是将最检索到的12个context重新排序后获取最相关的6个context。接下来我们就来通过该检索引擎来检索用户问题：

auto_merging_response = auto_merging_engine.query(
    "恐龙是冷血动物吗？"
)

这里由于我们在创建AutoMergingRetriever时设置了verbose=True因此检索器在检索相关文档时会显示检索过程的中间结果，从这些中间结果中我们看到其中有一个父节点中的3个叶子节点被检索到了，因为一个父节点包含最多4个叶子节点(由文档层次结构确定)，那么如果父节点中有3个叶子节点被检索到，那么该父节点将会作为context被返回给llm，而当只有1个叶子节点被检索到时，该父节点将不会被返回给llm。

下面我们看看LLM给出的最终答案：

from llama_index.response.notebook_utils import display_response

display_response(auto_merging_response)

这里我们提出了一个关于“恐龙为什么会灭绝？”问题，我们看到LLM给出了详细的解答。

下面我们给出完整的程序：

import os

from llama_index import (
    ServiceContext,
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)
from llama_index.node_parser import HierarchicalNodeParser
from llama_index.node_parser import get_leaf_nodes
from llama_index import StorageContext, load_index_from_storage
from llama_index.retrievers import AutoMergingRetriever
from llama_index.indices.postprocessor import SentenceTransformerRerank
from llama_index.query_engine import RetrieverQueryEngine
from llama_index.llms import OpenAI


def build_automerging_index(
    documents,
    llm,
    embed_model="local:BAAI/bge-small-zh-v1.5",
    save_dir="merging_index",
    chunk_sizes=None,
):
    chunk_sizes = chunk_sizes or [2048, 512, 128]
    node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes)
    nodes = node_parser.get_nodes_from_documents(documents)
    leaf_nodes = get_leaf_nodes(nodes)
    merging_context = ServiceContext.from_defaults(
        llm=llm,
        embed_model=embed_model,
    )
    storage_context = StorageContext.from_defaults()
    storage_context.docstore.add_documents(nodes)

    if not os.path.exists(save_dir):
        automerging_index = VectorStoreIndex(
            leaf_nodes, storage_context=storage_context, service_context=merging_context
        )
        automerging_index.storage_context.persist(persist_dir=save_dir)
    else:
        automerging_index = load_index_from_storage(
            StorageContext.from_defaults(persist_dir=save_dir),
            service_context=merging_context,
        )
    return automerging_index


def get_automerging_query_engine(
    automerging_index,
    similarity_top_k=12,
    rerank_top_n=6,
):
    base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k)
    retriever = AutoMergingRetriever(
        base_retriever, automerging_index.storage_context, verbose=True
    )
    rerank = SentenceTransformerRerank(
        top_n=rerank_top_n, model="BAAI/bge-reranker-base"
    )
    auto_merging_engine = RetrieverQueryEngine.from_args(
        retriever, node_postprocessors=[rerank]
    )
    return auto_merging_engine

index = build_automerging_index(
    [document],
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    save_dir="./merging_index",
)

query_engine = get_automerging_query_engine(index, similarity_top_k=6)

auto_merging_response = auto_merging_engine.query(
    "恐龙是冷血动物吗？"
)

display_response(auto_merging_response)

六、评估解释结果

在对自动合并检索进行评估时，我们需要先设置一组问题，然后我们让检索器来回答这组问题并生成答案，最后由评估器通过收集到的问题，答案，上下文(context)来综合评估LLM给出答案的质量，这里我们仍然使用TruLens三元组的评估指标(即Context Relevance，Groundedness和Answer Relevance的分数)来量化评估的效果,对此TruLens评估不熟悉朋友可以查看我之前写的TruLens 评估-扩大和加速LLM应用程序评估和评估句子-窗口检索这两篇博客中关于TruLens评估的介绍。下面我们做一下评估前的初始化工作：

from trulens_eval import Tru
from trulens_eval import Feedback,TruLlama
from trulens_eval import OpenAI as fOpenAI
from trulens_eval.feedback import Groundedness
import numpy as np


import nest_asyncio

#初始化评估数据库
Tru().reset_database()  

#设置线程的并发执行
nest_asyncio.apply()

接下来我们需要定义一组用来评估检索效果的问题：

#定义问题
eval_questions = [
                 "恐龙分哪几类?",
                 "体型最大的是哪种恐龙?",
                 "恐龙是怎么繁殖的？",
                 "恐龙是冷血动物吗？",
                 "恐龙为什么会灭绝? 什么时候灭绝的？",
                ]

然后会创建一个评估记录器函数，其中包含了选用的评估指标，最后我们会创建一个执行评估的函数：

#创建评估器对象
tru = Tru()
 
#定义评估记录器
def get_prebuilt_trulens_recorder(query_engine, app_id):
    openai = fOpenAI()
 
    qa_relevance = (
        Feedback(openai.relevance_with_cot_reasons, name="Answer Relevance")
        .on_input_output()
    )
 
    qs_relevance = (
        Feedback(openai.relevance_with_cot_reasons, name = "Context Relevance")
        .on_input()
        .on(TruLlama.select_source_nodes().node.text)
        .aggregate(np.mean)
    )
 
    grounded = Groundedness(groundedness_provider=openai)
 
    groundedness = (
        Feedback(grounded.groundedness_measure_with_cot_reasons, name="Groundedness")
            .on(TruLlama.select_source_nodes().node.text)
            .on_output()
            .aggregate(grounded.grounded_statements_aggregator)
    )
 
    feedbacks = [qa_relevance, qs_relevance, groundedness]
    tru_recorder = TruLlama(
        query_engine,
        app_id=app_id,
        feedbacks=feedbacks
    )
    return tru_recorder
 
#定义执行评估函数
def run_evals(eval_questions, tru_recorder, query_engine):
    for question in eval_questions:
        with tru_recorder as recording:
            response = query_engine.query(question)

二层结构评估

在完成了上述评估前的准备工作以后，接下来我们需要定义向量库索引，和创建检索引擎，最后执行检索，这里需要说明的是在本轮评估时我们采样的文档切割的层次结构为2层结构[2048,512]即叶子节点的块大小(chunk_size)为512，父节点的块大小为2048。

#创建向量库索引
auto_merging_index_0 = build_automerging_index(
    documents,
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    embed_model="local:BAAI/bge-small-zh-v1.5",
    save_dir="merging_index_0",
    chunk_sizes=[2048,512], 
)
#创建检索引擎
auto_merging_engine_0 = get_automerging_query_engine(
    auto_merging_index_0,
    similarity_top_k=12,
    rerank_top_n=6,
)
#创建记录器
tru_recorder = get_prebuilt_trulens_recorder(
    auto_merging_engine_0,
    app_id ='app_0'
)

#执行评估
run_evals(eval_questions, tru_recorder, auto_merging_engine_0)

当执行完评估完成以后，我们可以查看一下评估结果：

Tru().get_leaderboard(app_ids=[])

这里我们从本轮评估的结果中看到了三元组指标Context Relevance，Groundedness和Answer Relevance的分数，以及整个评估的耗时latency为15秒和评估的成本0.005433美刀，下面我们查看TruLens提供的基于streamlit的web app 界面:

Tru().run_dashboard()

三层结构评估

接下来我们将文档分割的层级结构设置为3层[2048,512,128], 这样我们的叶子节点的块大小将为128，而叶子节点的父文档的块大小则为512。

auto_merging_index_1 = build_automerging_index(
    documents,
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    embed_model="local:BAAI/bge-small-en-v1.5",
    save_dir="merging_index_1",
    chunk_sizes=[2048,512,128],
)

auto_merging_engine_1 = get_automerging_query_engine(
    auto_merging_index_1,
    similarity_top_k=12,
    rerank_top_n=6,
)

tru_recorder = get_prebuilt_trulens_recorder(
    auto_merging_engine_1,
    app_id ='app_1'
)

run_evals(eval_questions, tru_recorder, auto_merging_engine_1)

Tru().get_leaderboard(app_ids=[])

这里我们观察到当我们的文档切割方式采样了3层结构以后，三元组指标中Groundedness和Answer Relevance的分数都达到的最高值1，而Context Relevance略有下降，这是由于我们采用的是三层文档分割方式，此时叶子节点的父节点的块大小为512，而在之前二层结构切割文档时叶子节点的父节点的块大小为2048，因为我们每次只返回检索到的叶子节点的父节点文档，所以块大小为512的父文档的信息量要比块大小为2048的父文档的信息量要少，所以造成了Context Relevance的分数略有所下降，但是总成本total_cost为0.001654美刀，要比之前下降了一些，这是因为返回的父文档的token数只有之前的1/4，因此评估成本也相应的降下来了。下面我们查看TruLens提供的基于streamlit的web app 界面: