当前位置：首页 > news >正文

【RAG】【retrievers11】递归检索器 + 节点引用 + Braintrust评估

news 2026/5/23 16:53:13

案例目标

本案例展示如何使用递归检索(Recursive Retrieval)遍历节点关系，并基于"引用"获取节点。节点引用是一个强大的概念，在初次检索时，您可能希望获取引用而非原始文本。多个引用可以指向同一个节点。

案例探索了节点引用的不同用法：

分块引用：不同大小的分块引用更大的分块
元数据引用：摘要和生成的问题引用更大的分块

通过Braintrust评估系统，我们量化了递归检索+节点引用方法的效果，证明这种方法相比传统检索方法有显著提升。

技术栈与核心依赖

llama-index-llms-openai
llama-index-readers-file
llama-index-core
braintrust
autoevals
pypdf
transformers
torch

环境配置

# 安装必要的依赖
pip install llama-index-llms-openai llama-index-readers-file
pip install -U llama_hub llama_index braintrust autoevals pypdf pillow transformers torch torchvision
# 设置API密钥
import os

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"
os.environ["BRAINTRUST_API_KEY"] = "your_braintrust_api_key"
os.environ["TOKENIZERS_PARALLELISM"] = "true" # 避免Chroma的警告信息

案例实现

1. 数据准备

步骤 1

下载并加载Llama 2论文：

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

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

loader = PDFReader()
docs0 = loader.load_data(file=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)]

2. 创建基础节点

步骤 2

创建基础节点（分块大小1024）：

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import IndexNode

# 创建文本分割器
node_parser = SentenceSplitter(chunk_size=1024)

# 获取节点
base_nodes = node_parser.get_nodes_from_documents(docs)

# 设置节点ID
for idx, node in enumerate(base_nodes):
node.id_ = f"node-{idx}"

3. 基线检索器

步骤 3

创建基线检索器，通过嵌入相似度获取top-k原始文本节点：

from llama_index.core import VectorStoreIndex
from llama_index.core.embeddings import resolve_embed_model
from llama_index.llms.openai import OpenAI

# 设置嵌入模型和LLM
embed_model = resolve_embed_model("local:BAAI/bge-small-en")
llm = OpenAI(model="gpt-3.5-turbo")

# 创建向量索引和检索器
base_index = VectorStoreIndex(base_nodes, embed_model=embed_model)
base_retriever = base_index.as_retriever(similarity_top_k=2)

4. 分块引用：小子块引用大父块

步骤 4

构建小子块指向大父块的图结构：

from llama_index.core.retrievers import RecursiveRetriever

# 定义子块大小
sub_chunk_sizes = [128, 256, 512]
sub_node_parsers = [SentenceSplitter(chunk_size=c) for c in sub_chunk_sizes]

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)

步骤 5

创建递归检索器：

# 创建节点字典
all_nodes_dict = {n.node_id: n for n in all_nodes}

# 创建向量索引
vector_index_chunk = VectorStoreIndex(all_nodes, embed_model=embed_model)
vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=2)

# 创建递归检索器
retriever_chunk = RecursiveRetriever(
"vector",
retriever_dict={"vector": vector_retriever_chunk},
node_dict=all_nodes_dict,
verbose=True,
)

5. 元数据引用：摘要和生成的问题引用更大的块

步骤 6

提取元数据（摘要和问题）并创建引用：

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

# 创建提取器
extractors = [
SummaryExtractor(summaries=["self"], show_progress=True),
QuestionsAnsweredExtractor(questions=5, show_progress=True),
]

# 运行元数据提取器
metadata_dicts = []
for extractor in extractors:
metadata_dicts.extend(extractor.extract(base_nodes))

步骤 7

保存和加载元数据：

import json
import copy

# 保存元数据
def save_metadata_dicts(path):
with open(path, "w") as fp:
for m in metadata_dicts:
fp.write(json.dumps(m) + "\\n")

# 加载元数据
def load_metadata_dicts(path):
with open(path, "r") as fp:
metadata_dicts = [json.loads(l) for l in fp.readlines()]
return metadata_dicts

# 保存和加载
save_metadata_dicts("data/llama2_metadata_dicts.jsonl")
metadata_dicts = load_metadata_dicts("data/llama2_metadata_dicts.jsonl")

步骤 8

创建包含源节点和元数据的所有节点：

# 创建所有节点（源节点 + 元数据）
all_nodes = copy.deepcopy(base_nodes)
for idx, d in enumerate(metadata_dicts):
inode_q = IndexNode(
text=d["questions_this_excerpt_can_answer"],
index_id=base_nodes[idx].node_id,
)
inode_s = IndexNode(
text=d["section_summary"],
index_id=base_nodes[idx].node_id
)
all_nodes.extend([inode_q, inode_s])

# 创建节点字典
all_nodes_dict = {n.node_id: n for n in all_nodes}

# 创建向量索引和检索器
vector_index_metadata = VectorStoreIndex(all_nodes)
vector_retriever_metadata = vector_index_metadata.as_retriever(similarity_top_k=2)

# 创建递归检索器
retriever_metadata = RecursiveRetriever(
"vector",
retriever_dict={"vector": vector_retriever_metadata},
node_dict=all_nodes_dict,
verbose=True,
)

6. 评估设置

步骤 9

生成评估数据集：

from llama_index.core.evaluation import (
generate_question_context_pairs,
EmbeddingQAFinetuneDataset,
)
import nest_asyncio

nest_asyncio.apply()

# 生成问题-上下文对
eval_dataset = generate_question_context_pairs(base_nodes)
eval_dataset.save_json("data/llama2_eval_dataset.json")

# 加载数据集
eval_dataset = EmbeddingQAFinetuneDataset.from_json(
"data/llama2_eval_dataset.json"
)

步骤 10

定义评估指标和函数：

import pandas as pd
import braintrust

# 准备数据
queries = eval_dataset.queries
relevant_docs = eval_dataset.relevant_docs
data = [
({"input": queries[query], "expected": relevant_docs[query]})
for query in queries.keys()
]

# 定义评分函数
def hitRateScorer(input, expected, output=None):
is_hit = any([id in expected for id in output])
return 1 if is_hit else 0

def mrrScorer(input, expected, output=None):
for i, id in enumerate(output):
if id in expected:
return 1 / (i + 1)
return 0

步骤 11

评估分块检索器：

# 设置向量检索器相似度top k为更高值
top_k = 10

# 创建分块检索器
vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=10)
retriever_chunk = RecursiveRetriever(
"vector",
retriever_dict={"vector": vector_retriever_chunk},
node_dict=all_nodes_dict,
verbose=False,
)

# 定义运行函数
def runChunkRetriever(input, hooks):
retrieved_nodes = retriever_chunk.retrieve(input)
retrieved_ids = [node.node.node_id for node in retrieved_nodes]
return retrieved_ids

# 运行评估
chunkEval = await braintrust.Eval(
name="llamaindex-recurisve-retrievers",
data=data,
task=runChunkRetriever,
scores=[hitRateScorer, mrrScorer],
)

步骤 12

评估元数据检索器：

# 创建元数据检索器
vector_retriever_metadata = vector_index_metadata.as_retriever(similarity_top_k=10)
retriever_metadata = RecursiveRetriever(
"vector",
retriever_dict={"vector": vector_retriever_metadata},
node_dict=all_nodes_dict,
verbose=False,
)

# 定义运行函数
def runMetaDataRetriever(input, hooks):
retrieved_nodes = retriever_metadata.retrieve(input)
retrieved_ids = [node.node.node_id for node in retrieved_nodes]
return retrieved_ids

# 运行评估
metadataEval = await braintrust.Eval(
name="llamaindex-recurisve-retrievers",
data=data,
task=runMetaDataRetriever,
scores=[hitRateScorer, mrrScorer],
)

步骤 13

评估基线检索器：

# 创建基线检索器
base_retriever = base_index.as_retriever(similarity_top_k=10)

# 定义运行函数
def runBaseRetriever(input, hooks):
retrieved_nodes = base_retriever.retrieve(input)
retrieved_ids = [node.node.node_id for node in retrieved_nodes]
return retrieved_ids

# 运行评估
baseEval = await braintrust.Eval(
name="llamaindex-recurisve-retrievers",
data=data,
task=runBaseRetriever,
scores=[hitRateScorer, mrrScorer],
)