🚀 ONNX转换all-MiniLM-L6-v2
本项目是将sentence-transformers/all-MiniLM-L6-v2转换为ONNX模型。该模型可将句子和段落映射到384维的密集向量空间,适用于聚类或语义搜索等任务。此自定义模型会输出last_hidden_state
和pooler_output
,而使用默认ONNX配置导出的sentence-transformers模型仅包含last_hidden_state
作为输出。
🚀 快速开始
✨ 主要特性
📦 安装指南
若要使用此模型,需先安装optimum:
python -m pip install optimum
💻 使用示例
基础用法
from optimum.onnxruntime.modeling_ort import ORTModelForCustomTasks
from transformers import AutoTokenizer
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)
📚 详细文档
评估结果
若要对该模型进行自动评估,可参考 Sentence Embeddings Benchmark: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和Cloud团队成员在高效深度学习框架方面的支持。
预期用途
我们的模型旨在用作句子和短段落编码器。给定输入文本,它会输出一个捕获语义信息的向量。该句子向量可用于信息检索、聚类或句子相似度任务。默认情况下,输入文本超过256个词块时会被截断。
训练过程
预训练
我们使用了预训练的nreimers/MiniLM-L6-H384-uncased
模型。有关预训练过程的详细信息,请参考该模型的卡片。
微调
我们使用对比目标对模型进行微调。具体而言,我们计算批次中每个可能句子对的余弦相似度,然后通过与真实对进行比较来应用交叉熵损失。
超参数
我们在TPU v3 - 8上训练模型,共训练100k步,批次大小为1024(每个TPU核心128)。我们使用了500步的学习率预热,序列长度限制为128个标记。我们使用AdamW优化器,学习率为2e - 5。完整的训练脚本可在当前仓库中找到:train_script.py
。
训练数据
我们使用多个数据集的组合来微调模型,句子对总数超过10亿。我们根据加权概率对每个数据集进行采样,具体配置详见data_config.json
文件。
数据集 |
论文 |
训练元组数 |
Reddit comments (2015 - 2018) |
论文 |
726,484,430 |
S2ORC 引用对(摘要) |
论文 |
116,288,806 |
WikiAnswers 重复问题对 |
论文 |
77,427,422 |
PAQ (问题,答案)对 |
论文 |
64,371,441 |
S2ORC 引用对(标题) |
论文 |
52,603,982 |
S2ORC (标题,摘要) |
论文 |
41,769,185 |
Stack Exchange (标题,正文)对 |
- |
25,316,456 |
Stack Exchange (标题 + 正文,答案)对 |
- |
21,396,559 |
Stack Exchange (标题,答案)对 |
- |
21,396,559 |
MS MARCO 三元组 |
论文 |
9,144,553 |
GOOAQ: Open Question Answering with Diverse Answer Types |
论文 |
3,012,496 |
Yahoo Answers (标题,答案) |
论文 |
1,198,260 |
Code Search |
- |
1,151,414 |
COCO 图像描述 |
论文 |
828,395 |
SPECTER 引用三元组 |
论文 |
684,100 |
Yahoo Answers (问题,答案) |
论文 |
681,164 |
Yahoo Answers (标题,问题) |
论文 |
659,896 |
SearchQA |
论文 |
582,261 |
Eli5 |
论文 |
325,475 |
Flickr 30k |
论文 |
317,695 |
Stack Exchange 重复问题(标题) |
|
304,525 |
AllNLI (SNLI 和 MultiNLI) |
论文 SNLI,论文 MultiNLI |
277,230 |
Stack Exchange 重复问题(正文) |
|
250,519 |
Stack Exchange 重复问题(标题 + 正文) |
|
250,460 |
Sentence Compression |
论文 |
180,000 |
Wikihow |
论文 |
128,542 |
Altlex |
论文 |
112,696 |
Quora Question Triplets |
- |
103,663 |
Simple Wikipedia |
论文 |
102,225 |
Natural Questions (NQ) |
论文 |
100,231 |
SQuAD2.0 |
论文 |
87,599 |
TriviaQA |
- |
73,346 |
总计 |
|
1,170,060,424 |
📄 许可证
本项目采用Apache 2.0许可证。