🚀 Stella-PL-retrieval

这是一个基于 stella_en_1.5B_v5 的文本编码器，并针对波兰语信息检索任务进行了进一步微调。该编码器能够将文本转换为 1024 维的向量，专门针对波兰语信息检索任务进行了优化。

模型训练步骤

• 第一步：我们使用 2000 万条多样化的波兰语 - 英语文本对语料库，通过多语言知识蒸馏方法使模型适配波兰语。
• 第二步：使用由 140 万个查询组成的数据集，通过对比损失对模型进行微调。每个查询的正、负段落是借助 BAAI/bge-reranker-v2.5-gemma2-lightweight 重排器选择的。模型以 1024 个查询的批量大小训练了三个周期。

适用场景提示

如果您需要一个更通用的编码器，适用于语义相似度或聚类等更广泛的任务，建议使用第一步的蒸馏版本：sdadas/stella-pl。

🚀 快速开始

✨ 主要特性

• 基于 stella_en_1.5B_v5 进行微调，适用于波兰语信息检索任务。
• 能够将文本转换为 1024 维向量。
• 采用多语言知识蒸馏和对比损失微调，提升模型性能。

📦 安装指南

文档未提及具体安装步骤，可参考 sentence-transformers 进行安装。

💻 使用示例

基础用法

该模型使用与原始 stella_en_1.5B_v5 相同的提示。

• 对于检索任务，查询应添加前缀 "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "。
• 对于语义相似度等对称任务，两个文本都应添加前缀 "Instruct: Retrieve semantically similar text.\nQuery: "。

请注意，该模型使用自定义实现，因此在加载时应添加 trust_remote_code=True 参数。同时，建议使用 Flash Attention 2，可以通过 attn_implementation 参数启用。

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer(
    "sdadas/stella-pl-retrieval",
    trust_remote_code=True,
    device="cuda",
    model_kwargs={"attn_implementation": "flash_attention_2", "trust_remote_code": True}
)
model.bfloat16()

# 检索示例
query_prefix = "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "
queries = [query_prefix + "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."
]
queries_emb = model.encode(queries, convert_to_tensor=True, show_progress_bar=False)
answers_emb = model.encode(answers, convert_to_tensor=True, show_progress_bar=False)
best_answer = cos_sim(queries_emb, answers_emb).argmax().item()
print(answers[best_answer])

# 语义相似度示例
sim_prefix = "Instruct: Retrieve semantically similar text.\nQuery: "
sentences = [
    sim_prefix + "Trzeba zdrowo się odżywiać i uprawiać sport.",
    sim_prefix + "Warto jest prowadzić zdrowy tryb życia, uwzględniający aktywność fizyczną i dietę.",
    sim_prefix + "One should eat healthy and engage in sports.",
    sim_prefix + "Zakupy potwierdzasz PINem, który bezpiecznie ustalisz podczas aktywacji."
]
emb = model.encode(sentences, convert_to_tensor=True, show_progress_bar=False)
print(cos_sim(emb, emb))

📚 详细文档

评估结果

该模型在波兰语信息检索基准测试中实现了 NDCG@10 为 62.32 的成绩。详细结果请参阅 PIRB Leaderboard。

📄 许可证

此模型的许可证为 gemma。

🔧 引用信息

如果您使用该模型，请引用以下论文：

@article{dadas2024pirb,
  title={{PIRB}: A Comprehensive Benchmark of Polish Dense and Hybrid Text Retrieval Methods}, 
  author={Sławomir Dadas and Michał Perełkiewicz and Rafał Poświata},
  year={2024},
  eprint={2402.13350},
  archivePrefix={arXiv},
  primaryClass={cs.CL}
}