标签:

• mteb
• 句子转换器
• 句子相似度
• 句子转换器 语言:
• 英文 许可证: mit

E5-RoPE-基础版

LongEmbed: 扩展嵌入模型以支持长上下文检索。朱大伟、王亮、杨楠、宋一凡、吴文浩、韦福如、李素建，arxiv 2024。LongEmbed的GitHub仓库：https://github.com/dwzhu-pku/LongEmbed。

该模型包含12层，嵌入维度为768。

使用方法

以下是一个示例，展示如何对MS-MARCO段落排序数据集中的查询和段落进行编码。

import torch
import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
# 每个输入文本应以"query: "或"passage: "开头。
# 对于非检索任务，可以直接使用"query: "前缀。
input_texts = ['query: 女性每日应摄入多少蛋白质',
               'query: 峰会的定义',
               "passage: 作为一般指导原则，CDC建议19至70岁女性每日平均蛋白质摄入量为46克。但如本表所示，怀孕或备战马拉松时需增加摄入量。查看下表了解每日蛋白质建议摄入量。",
               "passage: 峰会定义（针对英语学习者）：1 山顶：山的最高点。2 最高级别。3 两国或多国政府领导人间的会议或系列会议。"]
tokenizer = AutoTokenizer.from_pretrained('dwzhu/e5rope-base', trust_remote_code=True)
model = AutoModel.from_pretrained('dwzhu/e5rope-base', trust_remote_code=True).cuda()
# 对输入文本进行分词
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)
batch_dict = {k: v.cuda() for k, v in batch_dict.items()}
outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
# 归一化嵌入向量
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())

训练细节

详情请参阅我们的论文：https://arxiv.org/abs/2404.12096.pdf。

基准评估

参考unilm/e5复现BEIR和MTEB基准上的评估结果。

请注意，E5-RoPE-基础版并非专为优化性能而训练。其目的是比较使用绝对位置嵌入（APE）和旋转位置嵌入（RoPE）的嵌入模型性能。通过对比E5-基础版和E5-RoPE-基础版，我们证明了基于RoPE的嵌入模型在有效处理长上下文方面的优势。详见我们的论文LongEmbed: 扩展嵌入模型以支持长上下文检索。

引用

如果您觉得我们的论文或模型有帮助，请考虑引用如下：

@article{zhu2024longembed,
  title={LongEmbed: Extending Embedding Models for Long Context Retrieval},
  author={Zhu, Dawei and Wang, Liang and Yang, Nan and Song, Yifan and Wu, Wenhao and Wei, Furu and Li, Sujian},
  journal={arXiv preprint arXiv:2404.12096},
  year={2024}
}