模型:
flax-sentence-embeddings/all_datasets_v3_mpnet-base
这是一个 sentence-transformers 模型:它将句子和段落映射到一个768维的稠密向量空间,可用于聚类或语义搜索等任务。
当您安装了 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)
如果没有 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 |