使用向量检索和rerank 在RAG数据集上实验评估hit_rate和mrr

背景

在前面部分 大模型生成RAG评估数据集并计算hit_rate 和 mrr 介绍了使用大模型生成RAG评估数据集与评估；

在 上文 使用到了BM25 关键词检索器。接下来，想利用向量检索器测试一下在RAG评估数据集上的 hit_rate 和 mrr；

简介

使用 向量检索 和 rerank 在给定RAG评估数据集上的实验计算 hit_rate 和 mrr；

对比了使用 rerank 和 不使用 rerank的实验结果；

步骤：

1. 基于RAG评估数据集，构建nodes节点；
2. 构建 CustomRetriever 自定义的检索器，在检索器中实现 向量检索和 rerank；
3. 实验评估；

代码实现

from typing import Listfrom llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.core.evaluation import RetrieverEvaluator
from llama_index.core.indices.postprocessor import SentenceTransformerRerank
from llama_index.core.indices.vector_store import VectorIndexRetriever
from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.settings import Settings
from llama_index.legacy.embeddings import HuggingFaceEmbedding
# from llama_index.legacy.schema import NodeWithScore, QueryBundle
from llama_index.core.schema import NodeWithScore, QueryBundle, QueryType, Node
from llama_index.core.evaluation import EmbeddingQAFinetuneDataset

利用数据集中的数据，构建nodes
pg_eval_dataset.json的下载地址: https://www.modelscope.cn/datasets/jieshenai/paul_graham_essay_rag/files

qa_dataset = EmbeddingQAFinetuneDataset.from_json("pg_eval_dataset.json")nodes = []
for key, value in qa_dataset.corpus.items():nodes.append(Node(id_=key, text=value))

m3e 向量编码模型
若想使用其他的编码模型，直接进行修改即可，modelscope和huggingface的编码模型都行；

from modelscope import snapshot_download
model_dir = snapshot_download('AI-ModelScope/m3e-base')
Settings.embed_model = HuggingFaceEmbedding(model_dir)
Settings.llm = None

由于huggingface被墙了，笔者使用的是 modelscope平台，model_dir 为编码模型在本地的绝对路径

自定义检索器

tok_k: 表示召回的节点数量，可自定义设置；

top_k = 10

定义向量检索器，还实现了rerank；

class CustomRetriever(BaseRetriever):"""Custom retriever that performs both Vector search and Knowledge Graph search"""def __init__(self, vector_retriever: VectorIndexRetriever, reranker=None) -> None:"""Init params."""super().__init__()self._vector_retriever = vector_retrieverself.reranker = rerankerdef _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:"""Retrieve nodes given query."""# print(query_bundle, isinstance(QueryBundle))retrieved_nodes = self._vector_retriever.retrieve(query_bundle)if self.reranker != 'None':retrieved_nodes = self.reranker.postprocess_nodes(retrieved_nodes, query_bundle)else:retrieved_nodes = retrieved_nodes[:top_k]return retrieved_nodesasync def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:"""Asynchronously retrieve nodes given query.Implemented by the user."""return self._retrieve(query_bundle)async def aretrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:if isinstance(str_or_query_bundle, str):str_or_query_bundle = QueryBundle(str_or_query_bundle)return await self._aretrieve(str_or_query_bundle)

eval_results包含每个query的 hit_rate 和 mrr，display_results 计算平均；

import pandas as pd
def display_results(eval_results):"""计算平均 hit_rate 和 mrr"""metric_dicts = []for eval_result in eval_results:metric_dict = eval_result.metric_vals_dictmetric_dicts.append(metric_dict)full_df = pd.DataFrame(metric_dicts)hit_rate = full_df["hit_rate"].mean()mrr = full_df["mrr"].mean()metric_df = pd.DataFrame({"hit_rate": [hit_rate], "mrr": [mrr]})return metric_df

向量检索实验

index = VectorStoreIndex(nodes)
vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=top_k)
retriever_evaluator = RetrieverEvaluator.from_metric_names(["mrr", "hit_rate"], retriever=vector_retriever
)
eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)
display_results(eval_results)

向量检索和rerank 实验

bge_reranker_base = SentenceTransformerRerank(model=snapshot_download("Xorbits/bge-reranker-base"),top_n=top_k)retriever = CustomRetriever(vector_retriever=vector_retriever,reranker=bge_reranker_base)retriever_evaluator = RetrieverEvaluator.from_metric_names(["mrr", "hit_rate"], retriever=retriever
)
eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)
display_results(eval_results)

若想使用其他的rerank模型，更换Xorbits/bge-reranker-base；

若使用modelscope平台的rerank模型，直接修改模型名即可；
若使用huggingface 平台的rerank模型，自行修改代码；

上述对比了，在向量检索下，对比了添加rerank和不添加rerank的实验结果；
如上图所示，相比只有向量检索的实验，加了rerank mrr 反而还下降了，这是一个比较反常的实验结果；

这个并不能说明rerank没有用，笔者在其他的RAG数据集测试时，rerank确实能提升mrr；本例子这里的情况大家忽略即可。
在本实验这里仅仅是给读者展示如何使用rerank；这也说明了rerank模型，也并不都能提升所有的mrr；

代码开源

本项目的完整代码，已发布到modelscope平台上；
点击下述链接查看代码：
https://www.modelscope.cn/datasets/jieshenai/paul_graham_essay_rag/file/view/master/vector_rerank_eval.ipynb?status=1