文章目录
- 背景
- 简介
- 代码实现
- 自定义检索器
- 向量检索实验
- 向量检索和rerank 实验
- 代码开源
背景
在前面部分 大模型生成RAG评估数据集并计算hit_rate 和 mrr 介绍了使用大模型生成RAG评估数据集与评估;
在 上文 使用到了BM25 关键词检索器。接下来,想利用向量检索器测试一下在RAG评估数据集上的 hit_rate 和 mrr;
简介
使用 向量检索 和 rerank 在给定RAG评估数据集上的实验计算 hit_rate 和 mrr;
对比了使用 rerank 和 不使用 rerank的实验结果;
步骤:
- 基于RAG评估数据集,构建nodes节点;
- 构建
CustomRetriever
自定义的检索器,在检索器中实现 向量检索和 rerank; - 实验评估;
代码实现
from typing import List
from 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_retriever
self.reranker = reranker
def _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_nodes
async 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_dict
metric_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