pipeline_tag: 句子相似度 tags:

• 句子转换器
• 特征提取
• 句子相似度 language: 英文 license: Apache-2.0许可证

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("我爱墨西哥卷饼！", return_tensors="pt")
pred = model(**inputs)

也可使用transformers的pipeline API：

from transformers import pipeline

onnx_extractor = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
text = "我爱墨西哥卷饼！"
pred = onnx_extractor(text)

评估结果

自动化评估结果参见句子嵌入基准测试：https://seb.sbert.net

背景

本项目旨在通过自监督对比学习目标，在超大规模句子级数据集上训练句子嵌入模型。我们基于预训练模型nreimers/MiniLM-L6-H384-uncased，在10亿句对数据集上进行微调。采用对比学习目标：给定句对中的一个句子，模型需从随机采样的其他句子中预测出实际配对的句子。 该模型开发于Hugging Face组织的使用JAX/Flax进行NLP与CV的社区周活动期间，属于用10亿训练对训练最佳句子嵌入模型项目。我们受益于高效硬件基础设施（7块TPU v3-8）以及Google Flax、JAX和云团队关于高效深度学习框架的指导。

用途

本模型用于句子和短段落编码，可将输入文本转换为蕴含语义信息的向量。所得句子向量可用于信息检索、聚类或句子相似度任务。默认截断超过256个word pieces的输入文本。

训练流程

预训练

使用预训练模型nreimers/MiniLM-L6-H384-uncased，详见模型卡获取预训练细节。

微调

采用对比目标进行微调：计算批次内所有可能句对的余弦相似度，通过对比真实句对应用交叉熵损失。

超参数

在TPU v3-8上训练10万步，批次大小1024（每TPU核心128）。学习率预热500步，序列长度限制为128个token。使用AdamW优化器，学习率2e-5。完整训练脚本见仓库：train_script.py。

训练数据

合并多个数据集进行微调，总句对数超10亿。各数据集采样权重配置见data_config.json文件。