ONNX转换为all-MiniLM-L6-v2

转换为 sentence-transformers/all-MiniLM-L6-v2

这是一个 sentence-transformers 的ONNX模型：它将句子和段落映射到一个384维的稠密向量空间，可用于聚类或语义搜索等任务。这个自定义模型接收last_hidden_state和pooler_output作为输入，而使用默认ONNX配置导出的sentence-transformers只包含last_hidden_state作为输出。

使用方法（HuggingFace Optimum）

当你安装了 optimum 之后，使用这个模型会变得很容易：

python -m pip install optimum

然后你可以像这样使用模型：

from optimum.onnxruntime.modeling_ort import ORTModelForCustomTasks

model = ORTModelForCustomTasks.from_pretrained("optimum/sbert-all-MiniLM-L6-with-pooler")
tokenizer = AutoTokenizer.from_pretrained("optimum/sbert-all-MiniLM-L6-with-pooler")
inputs = tokenizer("I love burritos!", return_tensors="pt")
pred = model(**inputs)

你还可以利用transformers中的pipeline API：

from transformers import pipeline

onnx_extractor = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
text = "I love burritos!"
pred = onnx_extractor(text)

评估结果

对于这个模型的自动评估结果，请参阅句子嵌入基准测试： https://seb.sbert.net

背景

该项目旨在使用自监督对比学习目标在非常大的句子级数据集上训练句子嵌入模型。我们使用预训练的 nreimers/MiniLM-L6-H384-uncased 模型，并在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和云团队成员的干预。

预期用途

我们的模型旨在用作句子和短段落编码器。给定一个输入文本，它输出一个能够捕捉语义信息的向量。该句向量可以用于信息检索、聚类或句子相似性任务。默认情况下，超过256个词块的输入文本将被截断。

训练过程

预训练

我们使用了预训练的 nreimers/MiniLM-L6-H384-uncased 模型。有关预训练过程的更多详细信息，请参阅模型卡片。

微调

我们使用对比目标对模型进行微调。具体而言，我们计算批处理中每对可能的句子对之间的余弦相似度。然后，通过与真实对进行比较，应用交叉熵损失。

我们在TPU v3-8上训练了模型。我们使用了批大小为1024（每个TPU核心128个）的训练步骤，共进行了100k步。我们使用了500个学习率预热。序列长度限制为128个标记。我们使用了学习率为2e-5的AdamW优化器。完整的训练脚本可在当前仓库的train_script.py中找到。

我们使用多个数据集的拼接来微调我们的模型。句子对的总数超过了10亿个句子。我们根据数据配置文件data_config.json中详细说明的加权概率对每个数据集进行采样。