英文

all-distilroberta-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-distilroberta-v1')
embeddings = model.encode(sentences)
print(embeddings)

用法(HuggingFace Transformers)

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

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-distilroberta-v1')
model = AutoModel.from_pretrained('sentence-transformers/all-distilroberta-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

背景

该项目旨在使用自监督的对比学习目标,在非常大的句级数据集上训练句子嵌入模型。我们使用预训练的 distilroberta-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团队成员的有关高效深度学习框架的干预。

预期用途

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

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

训练过程

预训练

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

微调

我们使用对比目标对模型进行微调。形式上,我们从批次中的每对句子计算余弦相似度。然后,将其与真实对进行比较,应用交叉熵损失。

超参数

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