许可证:apache-2.0
数据集:
- sentence-transformers/msmarco
语言:
- en
基础模型:
- nreimers/TinyBERT_L-6_H-768_v2
任务标签:文本排序
库名称:sentence-transformers
标签:
- transformers
MS Marco交叉编码器
该模型基于MS Marco段落排序任务训练而成。
适用于信息检索场景:给定查询语句时,可对所有候选段落(例如通过ElasticSearch检索获得)进行联合编码,并按相关性降序排序。详见SBERT.net检索与重排序指南。训练代码参见:SBERT.net训练MS Marco
Transformers调用方式
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
model = AutoModelForSequenceClassification.from_pretrained('cross-encoder/ms-marco-TinyBERT-L6')
tokenizer = AutoTokenizer.from_pretrained('cross-encoder/ms-marco-TinyBERT-L6')
features = tokenizer(['柏林有多少人口?', '柏林有多少人口?'],
['柏林注册居民达3,520,031人,占地面积891.82平方公里。', '纽约市以大都会艺术博物馆闻名。'],
padding=True, truncation=True, return_tensors="pt")
model.eval()
with torch.no_grad():
scores = model(**features).logits
print(scores)
SentenceTransformers调用方式
安装SentenceTransformers后可通过更简洁的方式调用:
from sentence_transformers import CrossEncoder
model = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L6', max_length=512)
scores = model.predict([('查询语句', '段落1'), ('查询语句', '段落2'), ('查询语句', '段落3')])
性能表现
下表列出不同预训练交叉编码器在TREC 2019深度学习赛道和MS Marco段落重排序数据集的表现:
模型名称 |
TREC DL19 NDCG@10 |
MS Marco Dev MRR@10 |
处理速度(篇/秒) |
V2系列模型 |
|
|
|
cross-encoder/ms-marco-TinyBERT-L2-v2 |
69.84 |
32.56 |
9000 |
cross-encoder/ms-marco-MiniLM-L2-v2 |
71.01 |
34.85 |
4100 |
cross-encoder/ms-marco-MiniLM-L4-v2 |
73.04 |
37.70 |
2500 |
cross-encoder/ms-marco-MiniLM-L6-v2 |
74.30 |
39.01 |
1800 |
cross-encoder/ms-marco-MiniLM-L12-v2 |
74.31 |
39.02 |
960 |
V1系列模型 |
|
|
|
cross-encoder/ms-marco-TinyBERT-L2 |
67.43 |
30.15 |
9000 |
cross-encoder/ms-marco-TinyBERT-L4 |
68.09 |
34.50 |
2900 |
cross-encoder/ms-marco-TinyBERT-L6 |
69.57 |
36.13 |
680 |
cross-encoder/ms-marco-electra-base |
71.99 |
36.41 |
340 |
其他模型 |
|
|
|
nboost/pt-tinybert-msmarco |
63.63 |
28.80 |
2900 |
nboost/pt-bert-base-uncased-msmarco |
70.94 |
34.75 |
340 |
nboost/pt-bert-large-msmarco |
73.36 |
36.48 |
100 |
Capreolus/electra-base-msmarco |
71.23 |
36.89 |
340 |
amberoad/bert-multilingual-passage-reranking-msmarco |
68.40 |
35.54 |
330 |
sebastian-hofstaetter/distilbert-cat-margin_mse-T2-msmarco |
72.82 |
37.88 |
720 |
注:速度测试基于V100 GPU。