e5-small-unsupervised - 无监督预训练模型用于提升文本嵌入与句子相似度

e5-small-unsupervised 项目介绍

项目概述

e5-small-unsupervised 是一个基于弱监督对比学习预训练的文本嵌入模型。这个模型与其前身 e5-small 很相似，但没有经过监督式微调。e5-small-unsupervised 模型包含了12个层，其嵌入尺寸为384。模型的设计旨在提高文本相似度的计算效率，适用于多个自然语言处理任务。

使用方法

使用 e5-small-unsupervised 模型时，每个输入文本需要以"query: "或"passage: "开头。以下是使用 MS-MARCO passage ranking 数据集对查询和段落进行编码的示例代码：

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, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day...",
    "passage: Definition of summit for English Language Learners. : 1 the highest point of a mountain : ..."
]

tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-small-unsupervised')
model = AutoModel.from_pretrained('intfloat/e5-small-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 基准测试上的评估结果。

支持 Sentence Transformers

e5-small-unsupervised 模型兼容于 Sentence Transformers 框架，以下示例展示了如何在此环境中使用该模型：

from sentence_transformers import SentenceTransformer

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