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

最新推荐文章于 2024-10-04 15:49:54 发布

jieshenai

最新推荐文章于 2024-10-04 15:49:54 发布

阅读量1.5k

点赞数 30

分类专栏：大模型文章标签：自然语言处理

本文链接：https://blog.csdn.net/sjxgghg/article/details/137412452

版权

大模型专栏收录该内容

21 篇文章 11 订阅

订阅专栏

文章目录

背景

在前面部分大模型生成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)