sbert-all-MiniLM-L6-with-pooler

项目概述

sbert-all-MiniLM-L6-with-pooler是一个高效的句子向量模型,它能够将句子和段落映射到384维的密集向量空间中。该模型基于sentence-transformers框架开发,适用于文本聚类、语义搜索等自然语言处理任务。

核心特点

• 该模型支持ONNX格式,可以轻松部署和使用
• 输出包含last_hidden_state和pooler_output两个特征
• 文本长度默认限制在256个词块
• 输出384维的句子向量,能够很好地捕捉语义信息
• 适用于信息检索、文本聚类和句子相似度计算等任务

技术实现

该模型基于nreimers/MiniLM-L6-H384-uncased预训练模型,并在超过10亿对句子数据上进行了微调。模型采用对比学习的方式训练,通过计算批次内所有可能的句子对之间的余弦相似度,使用交叉熵损失函数进行优化。

训练细节

• 使用TPU v3-8进行训练
• 训练步数为10万步
• 批次大小为1024
• 学习率预热步数为500
• 序列长度限制为128个token
• 使用AdamW优化器,学习率为2e-5

训练数据

模型的训练数据来自多个数据集,包括:

• Reddit评论数据(约7.2亿对)
• S2ORC学术文献数据(约2.1亿对)
• WikiAnswers问答数据(约7700万对)
• PAQ问答数据(约6400万对)
• Stack Exchange数据(约6800万对) 等多个高质量数据集,总计超过11.7亿对训练数据。

使用方法

模型可以通过Hugging Face的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")

应用场景

该模型适用于:

• 文本语义相似度计算
• 文档聚类分析
• 信息检索系统
• 智能问答系统
• 文本匹配任务