英文

all-mpnet-base-v1

这是一个 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-v1')
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-v1')
model = AutoModel.from_pretrained('sentence-transformers/all-mpnet-base-v1')

# 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 的一部分,我们从Google的Flax、JAX和Cloud团队成员的高效深度学习框架介入以及高效硬件基础设施(7个TPU v3-8)中受益。

预期用途

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

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

训练过程

预训练

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

微调

我们使用对比的目标对模型进行微调。形式上,我们从批处理中的每个可能的句对计算余弦相似度。然后,通过与真正的句对进行比较,应用交叉熵损失。

超参数

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

训练数据

我们使用了多个数据集的连接来微调我们的模型。句对的总数量超过了10亿个句子。我们根据数据_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
12324321 triplets 12325321 9,144,553
12326321 12327321 3,012,496
12328321 (Title, Answer) 12329321 1,198,260
12330321 - 1,151,414
12331321 Image captions 12332321 828,395
12333321 citation triplets 12334321 684,100
12328321 (Question, Answer) 12329321 681,164
12328321 (Title, Question) 12329321 659,896
12339321 12340321 582,261
12341321 12342321 325,475
12343321 12344321 317,695
12323321 Duplicate questions (titles) 304,525
AllNLI ( 12346321 and 12347321 12348321 , 12349321 277,230
12323321 Duplicate questions (bodies) 250,519
12323321 Duplicate questions (titles+bodies) 250,460
12352321 12353321 180,000
12354321 12355321 128,542
12356321 12357321 112,696
12358321 - 103,663
12359321 12360321 102,225
12361321 12362321 100,231
12363321 12364321 87,599
12365321 - 73,346
Total 1,124,818,467