49 LlamaIndex中的自动合并检索器：优化文本检索与合成

LlamaIndex中的自动合并检索器：优化文本检索与合成

在本文中，我们将介绍LlamaIndex中的AutoMergingRetriever，这是一种能够递归地“合并”引用父节点超过给定阈值的叶节点子集的检索器。这种方法允许我们将潜在不同的、较小的上下文合并成一个较大的上下文，从而有助于合成。

设置环境

首先，确保你已经安装了必要的库并设置了OpenAI API密钥：

import os
import openai

os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]

下载数据

下载示例数据，这里我们使用Llama 2论文作为测试数据：

!mkdir -p 'data/'
!wget --user-agent "Mozilla" "https://arxiv.org/pdf/2307.09288.pdf" -O "data/llama2.pdf"

加载数据

加载PDF文件并将其合并为一个文档：

from pathlib import Path
from llama_index.readers.file import PyMuPDFReader

loader = PyMuPDFReader()
docs0 = loader.load(file_path=Path("./data/llama2.pdf"))

from llama_index.core import Document

doc_text = "\n\n".join([d.get_content() for d in docs0])
docs = [Document(text=doc_text)]

解析文本并创建节点层次结构

使用HierarchicalNodeParser解析文本并创建节点层次结构：

from llama_index.core.node_parser import HierarchicalNodeParser, SentenceSplitter

node_parser = HierarchicalNodeParser.from_defaults()
nodes = node_parser.get_nodes_from_documents(docs)

from llama_index.core.node_parser import get_leaf_nodes, get_root_nodes

leaf_nodes = get_leaf_nodes(nodes)
root_nodes = get_root_nodes(nodes)

加载到存储中

定义文档存储并将所有节点加载到其中：

from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core import StorageContext
from llama_index.llms.openai import OpenAI

docstore = SimpleDocumentStore()
docstore.add_documents(nodes)

storage_context = StorageContext.from_defaults(docstore=docstore)

llm = OpenAI(model="gpt-3.5-turbo")

创建向量索引

将叶节点加载到向量索引中：

from llama_index.core import VectorStoreIndex

base_index = VectorStoreIndex(
    leaf_nodes,
    storage_context=storage_context,
)

定义检索器

创建AutoMergingRetriever并设置基础检索器：

from llama_index.core.retrievers import AutoMergingRetriever

base_retriever = base_index.as_retriever(similarity_top_k=6)
retriever = AutoMergingRetriever(base_retriever, storage_context, verbose=True)

查询与合并

执行查询并查看合并后的节点：

query_str = (
    "What could be the potential outcomes of adjusting the amount of safety"
    " data used in the RLHF stage?"
)

nodes = retriever.retrieve(query_str)
base_nodes = base_retriever.retrieve(query_str)

显示结果

显示合并后的节点和基础检索器的节点：

from llama_index.core.response.notebook_utils import display_source_node

for node in nodes:
    display_source_node(node, source_length=10000)

for node in base_nodes:
    display_source_node(node, source_length=10000)

插入查询引擎

将检索器插入查询引擎以合成自然语言响应：

from llama_index.core.query_engine import RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(retriever)
base_query_engine = RetrieverQueryEngine.from_args(base_retriever)

response = query_engine.query(query_str)
base_response = base_query_engine.query(query_str)

print(str(response))
print(str(base_response))

评估

评估层次检索器与基础检索器的性能：

from llama_index.core.evaluation import (
    CorrectnessEvaluator,
    SemanticSimilarityEvaluator,
    RelevancyEvaluator,
    FaithfulnessEvaluator,
    PairwiseComparisonEvaluator,
)

from llama_index.core.evaluation.eval_utils import (
    get_responses,
    get_results_df,
)
from llama_index.core.evaluation import BatchEvalRunner

eval_llm = OpenAI(model="gpt-4")

evaluator_c = CorrectnessEvaluator(llm=eval_llm)
evaluator_s = SemanticSimilarityEvaluator(llm=eval_llm)
evaluator_r = RelevancyEvaluator(llm=eval_llm)
evaluator_f = FaithfulnessEvaluator(llm=eval_llm)

evaluator_dict = {
    "correctness": evaluator_c,
    "faithfulness": evaluator_f,
    "relevancy": evaluator_r,
    "semantic_similarity": evaluator_s,
}

batch_runner = BatchEvalRunner(evaluator_dict, workers=2, show_progress=True)

eval_results = await batch_runner.aevaluate_responses(
    eval_qs, responses=pred_responses, reference=ref_response_strs
)
base_eval_results = await batch_runner.aevaluate_responses(
    eval_qs, responses=base_pred_responses, reference=ref_response_strs
)

results_df = get_results_df(
    [eval_results, base_eval_results],
    ["Auto Merging Retriever", "Base Retriever"],
    ["correctness", "relevancy", "faithfulness", "semantic_similarity"],
)

display(results_df)

分析

通过上述步骤，我们可以看到AutoMergingRetriever在合并节点方面的优势，它能够将较小的上下文合并成较大的上下文，从而有助于合成。通过评估，我们可以比较层次检索器与基础检索器的性能，从而更好地理解其优势和潜在的改进空间。

通过这种方法，LlamaIndex中的AutoMergingRetriever提供了一种强大的工具，用于优化文本检索和合成，从而提高检索结果的准确性和相关性。