英文

全MiniLM-L12-v1

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

用法(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-MiniLM-L12-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-MiniLM-L12-v1')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L12-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/MiniLM-L12-H384-uncased 模型,并在10亿个句对数据集上进行微调。我们使用对比学习目标:给定一句话,模型应该预测在我们的数据集中与其实际配对的一组随机采样的其他句子中的哪一个。

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

预期用途

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

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

训练过程

预训练

我们使用预训练的 microsoft/MiniLM-L12-H384-uncased 模型。请参阅模型卡片以获取有关预训练过程的更详细信息。

微调

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

超参数

我们在TPU v3-8上训练了我们的模型。我们使用批量大小为1024(每个TPU核心128个)进行了540k步的训练。我们使用了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
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