all-MiniLM-L12-v2

项目介绍: all-MiniLM-L12-v2

项目背景

all-MiniLM-L12-v2 是一个句子嵌入模型，隶属于SENTENCE-TRANSFORMERS 库。该模型的主要功能是将句子和段落映射到384维度的稠密向量空间，从而可以用于聚类或语义搜索等任务。

该项目的目标是利用大规模句子级数据集，通过自监督对比学习目标训练句子嵌入模型。我们采用了预训练的microsoft/MiniLM-L12-H384-uncased模型，并在十亿级别的句子对数据集上进行了微调。项目使用了 Hugging Face 组织的使用 JAX/Flax 的 NLP & CV 社区周中开发，得到了 Google Flax、JAX 和云团队的硬件支持。

预期用途

此模型主要用于对句子和短段落进行编码。给定一个输入文本，模型输出一个捕捉语义信息的向量，可用于信息检索、聚类或句子相似性任务。在默认配置下，超过256个词片段的输入文本会被截断。

模型使用

使用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-v2')
embeddings = model.encode(sentences)
print(embeddings)

使用HuggingFace Transformers库

不通过sentence-transformers，可以通过以下方式使用模型：

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

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    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 = ['This is an example sentence', 'Each sentence is converted']

tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L12-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L12-v2')

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

with torch.no_grad():
    model_output = model(**encoded_input)

sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)

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

训练过程

预训练

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

微调

通过对比学习目标微调模型。此过程中计算每个可能句子对的余弦相似性，并通过交叉熵损失与真实配对进行比较。

超参数

• 模型在 TPU v3-8 上训练，训练了100k步，批次大小为1024（每个TPU核心128）。
• 学习率的预热步为500，序列长度限制为128个词元。
• 使用AdamW优化器，学习率为2e-5。
• 完整的训练脚本在train_script.py中可访问。

训练数据

模型微调使用了多个数据集的合并，总句子对数量超过十亿。每个数据集的采样权重配置详见data_config.json文件。

总训练样本数量超过一十七亿条。

评估结果

有关该模型的自动化评估，请参见句子嵌入基准测试。