e5-base-unsupervised - E5-base突出文本嵌入的创新性

e5-base-unsupervised项目介绍

项目概述

e5-base-unsupervised 是一个基于预训练的文本嵌入模型，与 e5-base 类似，只是没有经过有监督的微调。该模型由12层神经网络组成，嵌入大小为768。它是在弱监督对比预训练的基础上，开发出的用于文本嵌入和句子相似度计算的强大工具。

模型使用方法

使用 e5-base-unsupervised 模型可以对来自 MS-MARCO 段落排名数据集的查询和段落进行编码。用户可以使用 Hugging Face 的 transformers 库来进行 token 化和编码。输入的文本以“query: ”或“passage: ”为前缀，这在执行检索等非对称任务时非常重要。而对于对称任务，如语义相似度计算和短语检索，使用“query: ”前缀即可。

使用示例代码如下：

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]

input_texts = [
    'query: how much protein should a female eat',
    'query: summit define',
    "passage: As a general guideline...",
    "passage: Definition of summit for English Language Learners..."
]

tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-base-unsupervised')
model = AutoModel.from_pretrained('intfloat/e5-base-unsupervised')

batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

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())

训练细节

该模型的训练细节参见原论文 Text Embeddings by Weakly-Supervised Contrastive Pre-training。该论文详细描述了模型的训练过程和使用的技术。

基准测试评价

用户可以参考 unilm/e5 来复现模型在 BEIR 和 MTEB benchmark 基准上的评价结果。

支持Sentence Transformers

e5-base-unsupervised 模型也可以与sentence_transformers库配合使用，用于编码文本的一致性和简便性。使用示例代码如下：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('intfloat/e5-base-unsupervised')
input_texts = [
    'query: how much protein should a female eat',
    'query: summit define',
    "passage: As a general guideline...",
    "passage: Definition of summit for English Language Learners..."
]
embeddings = model.encode(input_texts, normalize_embeddings=True)

安装所需包：

pip install sentence_transformers~=2.2.2

常见问题

1. 是否需要在输入文本前加上“query: ”和“passage: ”前缀？

是的，添加这些前缀是模型训练方式决定的，否则性能会下降。对于非对称任务，如开放性问答过程中的段落检索，使用“query: ”和“passage: ”；对于对称任务，如语义相似度或短语检索，仅使用“query: ”即可。

2. 为什么重现结果与模型卡中报告的结果略有不同？

这可能是因为使用了不同版本的 transformers 和 pytorch 库，造成了微小但不为零的性能差异。

限制

该模型仅支持英文文本，长文本将被截断至最多512个标记。

引用方式

如果您发现本文或模型有帮助，请引用如下：

@article{wang2022text,
  title={Text Embeddings by Weakly-Supervised Contrastive Pre-training},
  author={Wang, Liang and Yang, Nan and Huang, Xiaolong and Jiao, Binxing and Yang, Linjun and Jiang, Daxin and Majumder, Rangan and Wei, Furu},
  journal={arXiv preprint arXiv:2212.03533},
  year={2022}
}

通过以上信息，相信大家对e5-base-unsupervised 项目有了更为深入的了解。这是一款强大的文本嵌入工具，为自然语言处理的不同任务提供了有力的支持。