英文

all-mpnet-base-v2

这是一个 sentence-transformers 模型:它将句子和段落映射到一个768维的稠密向量空间,可用于聚类或语义搜索等任务。

使用(Sentence-Transformers)

当你安装了 sentence-transformers 后,使用这个模型变得很容易:

pip install -U sentence-transformers

然后你可以像这样使用模型:

from sentence_transformers import SentenceTransformer
sentences = ["This is an example sentence", "Each sentence is converted"]

model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
embeddings = model.encode(sentences)
print(embeddings)

使用(HuggingFace Transformers)

没有 sentence-transformers ,你可以像这样使用模型:首先,将输入通过变换器模型,然后在上下文词嵌入之上应用正确的汇集操作。

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F

#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)


# Sentences we want sentence embeddings for
sentences = ['This is an example sentence', 'Each sentence is converted']

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-mpnet-base-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-mpnet-base-v2')

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

# Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

# Perform pooling
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

# Normalize embeddings
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)

print("Sentence embeddings:")
print(sentence_embeddings)

评估结果

有关对该模型的自动评估,请参阅句子嵌入基准: https://seb.sbert.net

背景

该项目旨在使用自我监督的对比学习目标在非常大的句子级数据集上训练句子嵌入模型。我们使用预训练的 microsoft/mpnet-base 模型,并在10亿个句子对的数据集上进行微调。我们使用对比学习目标:给定一对句子中的一句话,模型应该预测在我们的数据集中,它实际上与哪个随机抽样的句子配对。

我们在由Hugging Face组织的 Community week using JAX/Flax for NLP & CV 期间开发了这个模型。我们作为项目中的一部分开发了这个模型: Train the Best Sentence Embedding Model Ever with 1B Training Pairs 。我们受益于高效的硬件基础设施来运行项目:7个TPU v3-8,以及来自Google的Flax、JAX和Cloud团队成员关于高效深度学习框架的干预。

预期用途

我们的模型旨在用作句子和短段落编码器。给定一个输入文本,它会输出一个捕捉语义信息的向量。该句向量可用于信息检索、聚类或句子相似性任务。

默认情况下,超过384个词片段的输入文本被截断。

训练过程

预训练

我们使用了预训练的 microsoft/mpnet-base 模型。有关预训练过程的更详细信息,请参阅模型卡。

微调

我们使用对比目标来微调模型。具体而言,我们从批次中的每个可能的句子对计算余弦相似度。然后将其与真实对比较应用交叉熵损失。

超参数

我们在TPU v3-8上训练了模型。我们使用大小为1024的批次大小进行了10万步的训练(每个TPU核心128个)。我们使用了500的学习率预热。序列长度限制为128个标记。我们使用了学习率为2e-5的AdamW优化器。完整的训练脚本可以在当前存储库中访问:train_script.py。

训练数据

我们使用多个数据集的串联来微调我们的模型。句子对的总数超过10亿个句子。我们根据数据配置文件(data_config.json)中详细说明的加权概率对每个数据集进行了抽样。

Dataset Paper Number of training tuples
12311321 12312321 726,484,430
12313321 Citation pairs (Abstracts) 12314321 116,288,806
12315321 Duplicate question pairs 12316321 77,427,422
12317321 (Question, Answer) pairs 12318321 64,371,441
12313321 Citation pairs (Titles) 12314321 52,603,982
12313321 (Title, Abstract) 12314321 41,769,185
12323321 (Title, Body) pairs - 25,316,456
12323321 (Title+Body, Answer) pairs - 21,396,559
12323321 (Title, Answer) pairs - 21,396,559
12326321 triplets 12327321 9,144,553
12328321 12329321 3,012,496
12330321 (Title, Answer) 12331321 1,198,260
12332321 - 1,151,414
12333321 Image captions 12334321 828,395
12335321 citation triplets 12336321 684,100
12330321 (Question, Answer) 12331321 681,164
12330321 (Title, Question) 12331321 659,896
12341321 12342321 582,261
12343321 12344321 325,475
12345321 12346321 317,695
12323321 Duplicate questions (titles) 304,525
AllNLI ( 12348321 and 12349321 12350321 , 12351321 277,230
12323321 Duplicate questions (bodies) 250,519
12323321 Duplicate questions (titles+bodies) 250,460
12354321 12355321 180,000
12356321 12357321 128,542
12358321 12359321 112,696
12360321 - 103,663
12361321 12362321 102,225
12363321 12364321 100,231
12365321 12366321 87,599
12367321 - 73,346
Total 1,170,060,424