RAG 评测数据集建设尚处于初期阶段，缺乏针对特定领域和场景的专业数据集。市面上常见的 MS-Marco 和 BEIR 数据集覆盖范围有限，且在实际使用场景中效果可能与评测表现不符。目前最权威的检索榜单是 HuggingFace MTEB，今天我们来学习使用MTEB，并来评测自研模型recall效果。
MTEB 是一个包含广泛文本嵌入（Text Embedding）的基准测试，它提供了多种语言的数十个数据集，用于各种 NLP 任务，例如文本分类、聚类、检索和文本相似性。MTEB 提供了一个公共排行榜，允许研究人员提交他们的结果并跟踪他们的进展。MTEB 还提供了一个简单的 API，允许研究人员轻松地将他们的模型与基准测试进行比较。
安装使用
pip install mteb
使用入门
- 最简单的用法就是，直接编写python代码来测试 (see scripts/run_mteb_english.py and mteb/mtebscripts for more):
from mteb import MTEB
from sentence_transformers import SentenceTransformer
# Define the sentence-transformers model name
model_name = "average_word_embeddings_komninos"
model = SentenceTransformer(model_name)
evaluation = MTEB(tasks=["Banking77Classification"])
results = evaluation.run(model, output_folder=f"results/{model_name}")
- 也可以使用官方提供的 CLI
mteb --available_tasks
mteb -m average_word_embeddings_komninos \
-t Banking77Classification \
--output_folder results/average_word_embeddings_komninos \
--verbosity 3
高级用法
测试数据集选择
MTEB支持指定数据集，可以通过下面的形式
- 按task_type任务类型（例如“聚类”或“分类”)
evaluation = MTEB(task_types=['Clustering', 'Retrieval']) # Only select clustering and retrieval tasks
- 按类别划分, 例如“句子到句子 "S2S" (sentence to sentence) "P2P" (paragraph to paragraph)
evaluation = MTEB(task_categories=['S2S']) # Only select sentence2sentence datasets
- 按照文本语言
evaluation = MTEB(task_langs=["en", "de"]) # Only select datasets which are "en", "de" or "en-de"
还可以针对数据集选择语言：
from mteb.tasks import AmazonReviewsClassification, BUCCBitextMining
evaluation = MTEB(tasks=[
AmazonReviewsClassification(langs=["en", "fr"]) # Only load "en" and "fr" subsets of Amazon Reviews
BUCCBitextMining(langs=["de-en"]), # Only load "de-en" subset of BUCC
])
可为某些任务集合提供预设
from mteb import MTEB_MAIN_EN
evaluation = MTEB(tasks=MTEB_MAIN_EN, task_langs=["en"])
自定义评测 split
有的数据集有多个split，评测会比较耗时，可以指定splits，来减少评测时间，比如下面的就指定了只用test split。
evaluation.run(model, eval_splits=["test"])
自定义评测模型
如果想自定义评测模型，可以自定义一个类，只要实现一个encode函数，输入是一个句子列表，返回的是一个嵌入向量列表（嵌入可以是np.array、torch.tensor等）。可以参考 mteb/mtebscripts repo 仓库。
class MyModel():
def encode(self, sentences, batch_size=32, **kwargs):
"""
Returns a list of embeddings for the given sentences.
Args:
sentences (`List[str]`): List of sentences to encode
batch_size (`int`): Batch size for the encoding
Returns:
`List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences
"""
pass
model = MyModel()
evaluation = MTEB(tasks=["Banking77Classification"])
evaluation.run(model)
如果针对query和corpus需要使用不同的encode方法，可以独立提供encode_queries and encode_corpus两个方法。
class MyModel():
def encode_queries(self, queries, batch_size=32, **kwargs):
"""
Returns a list of embeddings for the given sentences.
Args:
queries (`List[str]`): List of sentences to encode
batch_size (`int`): Batch size for the encoding
Returns:
`List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences
"""
pass
def encode_corpus(self, corpus, batch_size=32, **kwargs):
"""
Returns a list of embeddings for the given sentences.
Args:
corpus (`List[str]` or `List[Dict[str, str]]`): List of sentences to encode
or list of dictionaries with keys "title" and "text"
batch_size (`int`): Batch size for the encoding
Returns:
`List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences
"""
pass
自定义评测Task（数据集）
要添加一个新任务，你需要实现一个从与任务类型相关的AbsTask继承的新类（例如，对于重排任务是AbsTaskReranking）。你可以在这里找到支持的任务类型。
比如下面的自定义重排任务：
from mteb import MTEB
from mteb.abstasks.AbsTaskReranking import AbsTaskReranking
from sentence_transformers import SentenceTransformer
class MindSmallReranking(AbsTaskReranking):
@property
def description(self):
return {
"name": "MindSmallReranking",
"hf_hub_name": "mteb/mind_small",
"description": "Microsoft News Dataset: A Large-Scale English Dataset for News Recommendation Research",
"reference": "https://www.microsoft.com/en-us/research/uploads/prod/2019/03/nl4se18LinkSO.pdf",
"type": "Reranking",
"category": "s2s",
"eval_splits": ["validation"],
"eval_langs": ["en"],
"main_score": "map",
}
model = SentenceTransformer("average_word_embeddings_komninos")
evaluation = MTEB(tasks=[MindSmallReranking()])
evaluation.run(model)
源码分析
Retrieval召回评测
召回评测是通过RetrievalEvaluator类实现的。
def __init__(
self,
queries: Dict[str, str], # qid => query
corpus: Dict[str, str], # cid => doc
relevant_docs: Dict[str, Set[str]], # qid => Set[cid]
corpus_chunk_size: int = 50000,
mrr_at_k: List[int] = [10],
ndcg_at_k: List[int] = [10],
accuracy_at_k: List[int] = [1, 3, 5, 10],
precision_recall_at_k: List[int] = [1, 3, 5, 10],
map_at_k: List[int] = [100],
show_progress_bar: bool = False,
batch_size: int = 32,
name: str = "",
score_functions: List[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {
"cos_sim": cos_sim,
"dot": dot_score,
}, # Score function, higher=more similar
main_score_function: str = None,
limit: int = None,
**kwargs
):
super().__init__(**kwargs)
self.queries_ids = []
for qid in queries:
if qid in relevant_docs and len(relevant_docs[qid]) > 0:
self.queries_ids.append(qid)
if limit and len(self.queries_ids) >= limit:
break
self.queries = [queries[qid] for qid in self.queries_ids]
self.corpus_ids = list(corpus.keys())
self.corpus = [corpus[cid] for cid in self.corpus_ids]
self.relevant_docs = relevant_docs
self.corpus_chunk_size = corpus_chunk_size
self.mrr_at_k = mrr_at_k
self.ndcg_at_k = ndcg_at_k
self.accuracy_at_k = accuracy_at_k
self.precision_recall_at_k = precision_recall_at_k
self.map_at_k = map_at_k
self.show_progress_bar = show_progress_bar
self.batch_size = batch_size
self.name = name
self.score_functions = score_functions
self.score_function_names = sorted(list(self.score_functions.keys()))
self.main_score_function = main_score_function
构造函数几个重要的参数：
- queries: Dict[str, str], # qid => query qid到query的dict
- corpus: Dict[str, str], # cid => doc docid到doc的dict
- relevant_docs: Dict[str, Set[str]], # qid => Set[cid] qid到相关docid的dict
因此，要自定义评测任务，需要提供这些数据。
具体的评测函数在compute_metrics里：
def compute_metrics(self, model, corpus_model=None, corpus_embeddings: torch.Tensor = None) -> Dict[str, float]:
if corpus_model is None:
corpus_model = model
max_k = max(
max(self.mrr_at_k),
max(self.ndcg_at_k),
max(self.accuracy_at_k),
max(self.precision_recall_at_k),
max(self.map_at_k),
)
# Compute embedding for the queries
logger.info("Encoding the queries...")
# We don't know if encode has the kwargs show_progress_bar
kwargs = {
"show_progress_bar": self.show_progress_bar
} if "show_progress_bar" in inspect.signature(model.encode).parameters else {}
query_embeddings = np.asarray(model.encode(self.queries, batch_size=self.batch_size, **kwargs))
queries_result_list = {}
for name in self.score_functions:
queries_result_list[name] = [[] for _ in range(len(query_embeddings))]
# Iterate over chunks of the corpus
logger.info("Encoding chunks of corpus, and computing similarity scores with queries...")
for corpus_start_idx in trange(
0,
len(self.corpus),
self.corpus_chunk_size,
desc="Corpus Chunks",
disable=not self.show_progress_bar,
):
# Encode chunk of corpus
if corpus_embeddings is None:
corpus_end_idx = min(corpus_start_idx + self.corpus_chunk_size, len(self.corpus))
sub_corpus_embeddings = np.asarray(corpus_model.encode(
self.corpus[corpus_start_idx:corpus_end_idx],
batch_size=self.batch_size,
))
else:
corpus_end_idx = min(corpus_start_idx + self.corpus_chunk_size, len(corpus_embeddings))
sub_corpus_embeddings = corpus_embeddings[corpus_start_idx:corpus_end_idx]
# Compute cosine similarites
for name, score_function in self.score_functions.items():
pair_scores = score_function(query_embeddings, sub_corpus_embeddings)
# Get top-k values
pair_scores_top_k_values, pair_scores_top_k_idx = torch.topk(
pair_scores,
min(max_k, len(pair_scores[0])),
dim=1,
largest=True,
sorted=False,
)
pair_scores_top_k_values = pair_scores_top_k_values.cpu().tolist()
pair_scores_top_k_idx = pair_scores_top_k_idx.cpu().tolist()
for query_itr in range(len(query_embeddings)):
for sub_corpus_id, score in zip(
pair_scores_top_k_idx[query_itr],
pair_scores_top_k_values[query_itr],
):
corpus_id = self.corpus_ids[corpus_start_idx + sub_corpus_id]
queries_result_list[name][query_itr].append({"corpus_id": corpus_id, "score": score})
# Compute scores
logger.info("Computing metrics...")
scores = {name: self._compute_metrics(queries_result_list[name]) for name in self.score_functions}
return scores
model（embedding模型），corpus_model（如果doc用单独的embedding模型，需要传入这个参数，否则默认使用和query一样的model）- 首先会计算query_embedding
query_embeddings = np.asarray(model.encode(self.queries, batch_size=self.batch_size, **kwargs)) - 然后计算corpus_embeddings
- 通过score_function，计算tok_k， 结果放到queries_result_list
- 根据召回结果计算指标
_compute_metrics, 会计算"mrr@k", "ndcg@k", "accuracy@k", "precision_recall@k", "map@k"等指标
Reranking 精排
精排是通过RerankingEvaluator来实现的。
class RerankingEvaluator(Evaluator):
"""
This class evaluates a SentenceTransformer model for the task of re-ranking.
Given a query and a list of documents, it computes the score [query, doc_i] for all possible
documents and sorts them in decreasing order. Then, MRR@10 and MAP is compute to measure the quality of the ranking.
:param samples: Must be a list and each element is of the form:
- {'query': '', 'positive': [], 'negative': []}. Query is the search query, positive is a list of positive
(relevant) documents, negative is a list of negative (irrelevant) documents.
- {'query': [], 'positive': [], 'negative': []}. Where query is a list of strings, which embeddings we average
to get the query embedding.
"""
def __init__(
self,
samples,
mrr_at_k: int = 10,
name: str = "",
similarity_fct=cos_sim,
batch_size: int = 512,
use_batched_encoding: bool = True,
limit: int = None,
**kwargs,
):
给定一个query和一组文档，模型计算文档得分，并按降序排列，最后计算MRR@10和MAP指标来衡量排名的质量。
__init__方法接收以下参数：
samples：必须是一个列表，每个元素的形式为：- {'query': '', 'positive': [], 'negative': []}。查询是搜索查询，正文档是相关（正面）文档的列表，负文档是无关（负面）文档的列表。
- {'query': [], 'positive': [], 'negative': []}。其中查询是一个字符串列表，我们将这些字符串的平均嵌入作为查询嵌入。
mrr_at_k：默认值为10，表示计算MRR时考虑的前k个结果。name：默认值为空字符串，表示评估器的名称。similarity_fct：默认值为cos_sim，表示用于计算相似度的函数。
在compute_metrics_batched 计算得分，还是计算的cos得分，这里相当于直接计算的embedding的排序能力，如果要计算cross模型的排序能力，默认的代码不适用，需要重新定制。
评测实践
说了这么多，现在切入正题：
- 评测自研模型的召回能力 —— 自定义模型
- 自定义评测集，对比开源模型和自研模型的效果 —— 自定义评测任务
自研模型召回效果评测
我们先评估模型召回效果，训练好的模型导出为onnx，因此我们通过onnxrutime来进行推理，先自定义模型：
from mteb import MTEB
import onnxruntime as ort
from paddlenlp.transformers import AutoTokenizer
import math
from tqdm import tqdm
# 模型路径
model_path = "onnx/fp16_model.onnx"
tokenizer_path = "model_520000"
class MyModel():
def __init__(self, use_gpu=True):
providers = ['CUDAExecutionProvider'] if use_gpu else ['CPUExecutionProvider']
sess_options = ort.SessionOptions()
self.predictor = ort.InferenceSession(
model_path, sess_options=sess_options, providers=providers)
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
def encode(self, sentences, batch_size=64, **kwargs):
all_embeddings = []
# 向上取整
batch_count = math.ceil(len(sentences) / batch_size)
for i in tqdm(range(batch_count)):
# 按batch
sub_sentences = sentences[i * batch_size : min(len(sentences), (i + 1) * batch_size)]
features = self.tokenizer(sub_sentences, max_seq_len=128,
pad_to_max_seq_len=True, truncation_strategy="longest_first")
vecs = self.predictor.run(None, features.data)
all_embeddings.extend(vecs[0])
return all_embeddings
由于传进来的sentences是所有的数据，我们需要按照batch_size，分批进行embedding计算，计算好的放入all_embeddings，最后返回即可。
自定义召回评测任务
上面分析源代码时提到了，自定义时需要提供qurey，doc，以及query的相关doc
假设我们的自定义测试为jsonline格式，每行包含query，以及相关的doc，json格式如下：
{
"query": "《1984》是什么",
"data": [
{
"title": "《1984》介绍-知乎",
"summary": "《1984》是伪装成小说的政治思想...",
"url": "",
"id": 5031622209044687985,
"answer": "完全相关",
"accuracy": "无错",
"result": "good"
}
]
}
那么我们可以编写自定义召回评测任务：
class SSRetrieval(AbsTaskRetrieval):
@property
def description(self):
return {
'name': 'SSRetrieval',
'description': 'SSRetrieval是S研发部测试团队准备的召回测试集',
'type': 'Retrieval',
'category': 's2p',
'json_path': '/data/xapian-core-1.4.24/demo/result.json',
'eval_splits': ['dev'],
'eval_langs': ['zh'],
'main_score': 'recall_at_10',
}
def load_data(self, **kwargs):
if self.data_loaded:
return
self.corpus = {} # doc_id => doc
self.queries = {} # qid => query
self.relevant_docs = {} # qid => Set[doc_id]
query_index = 1
with open(self.description['json_path'], 'r', encoding='utf-8') as f:
for line in f:
if "完全相关" not in line:
continue
line = json.loads(line)
query = line['query']
query_id = str(query_index)
self.queries[query_id] = query
query_index = query_index + 1
query_relevant_docs = []
for doc in line['data']:
doc_id = str(doc['id'])
self.corpus[doc_id] = {"title": doc["title"], "text": doc["summary"]}
if doc['answer'] == "完全相关":
if query_id not in self.relevant_docs:
self.relevant_docs[query_id] = {}
self.relevant_docs[query_id][doc_id] = 1
# debug使用
# if query_index == 100:
# break
self.queries = DatasetDict({"dev": self.queries})
self.corpus = DatasetDict({"dev": self.corpus})
self.relevant_docs = DatasetDict({"dev": self.relevant_docs})
self.data_loaded = True
用自定义模型，评测自定义任务
if __name__ == '__main__':
model = MyModel()
# task_names = [t.description["name"] for t in MTEB(task_types='Retrieval',
# task_langs=['zh', 'zh-CN']).tasks]
task_names = ["SSRetrieval"]
for task in task_names:
model.query_instruction_for_retrieval = None
evaluation = MTEB(tasks=[task], task_langs=['zh', 'zh-CN'])
evaluation.run(model, output_folder=f"zh_results/256_model", batch_size=64)
总结
mteb 最为embedding召回效果测试，是一个权威的榜单，本身提供的工具框架也具备较好的扩展性，方便开发者自定义模型和自定义评测任务。
