LlamaIndex中的自动合并检索：原理与实践

LlamaIndex中的自动合并检索：原理与实践

引用

CSDN

1.

https://blog.csdn.net/weixin_42608414/article/details/135666037

自动合并检索（Auto-merging Retrieval）是LlamaIndex中的一种高级RAG（Retrieval-Augmented Generation）技术，它通过将文档按特定的层次结构进行切割和检索，能够更有效地组织和利用文档信息。本文将详细介绍自动合并检索的工作原理、实现步骤以及评估方法，帮助读者更好地理解和应用这一技术。

一、文档的层次结构

文档层次结构是指在切割文档时按照特定的层次结构进行，可以理解为广度和深度两个维度。具体来说：

• 广度（兄弟关系）：指相同层级的文档块之间的关系，例如章节之间、段落之间、句子之间的关系。
• 深度（父子关系）：指不同层级的文档块之间的关系，例如章节和段落之间、段落和句子之间的关系。

为了更好地理解文档层次结构，可以参考下图所示的DocArray项目示意图：

二、自动合并检索

LlamaIndex的自动合并检索首先需要将文档按特定的层次结构进行切割。例如，定义切割的层次结构为[1024, 512, 128]，表示将文档按三层结构进行切割，顶层节点的块大小为1024，中间层的块大小为512，底层的叶子节点的块大小为128。在检索时，只拿叶子节点和问题进行匹配，当某个父节点下的多数叶子节点都与问题匹配上，则将该父节点作为context返回给LLM。

实战环节

1. 环境配置

首先需要安装以下Python包：

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

然后导入OpenAI的API密钥：

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

2. 加载数据

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

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

3. 设置文档层次结构

创建一个文档切割器，按指定的文档层次结构切割文档：

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

使用文档切割器切割下载的科普文章：

nodes = node_parser.get_nodes_from_documents(documents)
len(nodes)  # 查看切割后的文档块数量

查看某个文档块的内容：

nodes[50]

查看叶子节点文档：

from llama_index.node_parser import get_leaf_nodes
leaf_nodes = get_leaf_nodes(nodes)
leaf_nodes[30]

查看叶子节点所属的父文档信息：

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

4. 创建向量库索引

创建向量库索引，使用OpenAI的gpt-3.5-turbo作为LLM，bge-small-zh-v1.5作为embedding模型：

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

llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
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")

5. 定义和执行检索器

创建自动合并检索器和检索引擎：

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])

通过检索引擎检索用户问题：

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

查看LLM给出的最终答案：

from llama_index.response.notebook_utils import display_response
display_response(auto_merging_response)

6. 评估解释结果

使用TruLens进行评估，定义评估问题：

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

创建评估记录器和执行评估：

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()

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)

# 二层结构评估
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=[])

# 三层结构评估
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=[])

评估结果显示，采用三层文档分割方式时，Groundedness和Answer Relevance的分数达到最高值1，但Context Relevance略有下降。这是因为叶子节点的父文档的块大小为512，而二层结构切割时为2048，导致信息量减少。然而，总成本有所下降，因为返回的父文档的token数只有之前的1/4。

总结

自动合并检索是对句子-窗口检索的补充，通过调整文档层次结构参数，可以优化检索效果和成本。建议使用不同的文档层次结构参数进行迭代，使用RAG三元组评估应用程序版本，跟踪实验以选择最佳的文档分割层次结构参数。