英文

all-roberta-large-v1

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

使用方法(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-roberta-large-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-roberta-large-v1')
model = AutoModel.from_pretrained('sentence-transformers/all-roberta-large-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

背景

该项目旨在使用自监督的对比学习目标,在非常大的句子级数据集上训练句子嵌入模型。我们使用预训练的 roberta-large 模型,并在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,以及来自谷歌的Flax,JAX和Cloud团队成员关于高效深度学习框架的干预。

预期用途

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

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

训练过程

预训练

我们使用预训练的 roberta-large 。有关预训练过程的更详细信息,请参考模型卡。

微调

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

网络参数

我们在TPU v3-8上训练了模型。我们使用256的批量大小(每个TPU核心32个)。我们使用了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