🚀 波兰重排器RoBERTa v2

这是一个基于 sdadas/polish-roberta-large-v2 的重排器改进版本，它在一个大型文本对数据集上使用 RankNet损失 进行训练。该模型的训练方式和使用的数据与 sdadas/polish-reranker-large-ranknet 相同，但有以下改进：

• 我们使用 BAAI/bge-reranker-v2.5-gemma2-lightweight 的预测结果进行蒸馏，而非 unicamp-dl/mt5-13b-mmarco-100k。
• 我们使用了支持Flash Attention 2的RoBERTa模型的自定义实现。如果您想使用这些特性，请使用参数 trust_remote_code=True 和 attn_implementation="flash_attention_2" 加载模型。

我们的重排器在PIRB基准测试中取得了接近 BAAI/bge-reranker-v2.5-gemma2-lightweight 的结果，甚至在某些数据集上超越了它。同时，它的参数规模比后者小21倍多 —— 分别为4.35亿和92.4亿。

🚀 快速开始

本模型可通过Huggingface Transformers库使用，以下是具体的使用方法。

💻 使用示例

基础用法

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np

query = "Jak dożyć 100 lat?"
answers = [
    "Trzeba zdrowo się odżywiać i uprawiać sport.",
    "Trzeba pić alkohol, imprezować i jeździć szybkimi autami.",
    "Gdy trwała kampania politycy zapewniali, że rozprawią się z zakazem niedzielnego handlu."
]

model_name = "sdadas/polish-reranker-roberta-v2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="cuda"
)
texts = [f"{query}</s></s>{answer}" for answer in answers]
tokens = tokenizer(texts, padding="longest", max_length=512, truncation=True, return_tensors="pt").to("cuda")
output = model(**tokens)
results = output.logits.detach().cpu().float().numpy()
results = np.squeeze(results)
print(results.tolist())

📚 详细文档

评估结果

该模型在波兰信息检索基准测试的重排器类别中实现了 NDCG@10 为 65.30 的成绩。详细结果请参考 PIRB排行榜。

引用信息

@article{dadas2024assessing,
  title={Assessing generalization capability of text ranking models in Polish}, 
  author={Sławomir Dadas and Małgorzata Grębowiec},
  year={2024},
  eprint={2402.14318},
  archivePrefix={arXiv},
  primaryClass={cs.CL}
}

📄 许可证

本项目采用gemma许可证。

📦 模型信息