🚀 句子相似度模型

本项目旨在使用自监督对比学习目标，在超大型句子级数据集上训练句子嵌入模型。通过该模型，能够将输入的句子转化为向量，捕捉句子的语义信息，可广泛应用于信息检索、聚类或句子相似度任务等领域。

🚀 快速开始

本模型可作为句子编码器使用。给定输入句子，它将输出一个捕获句子语义信息的向量。句子向量可用于信息检索、聚类或句子相似度任务。

✨ 主要特性

• 大规模数据训练：使用超过10亿个句子对的数据集进行微调，使模型能学习到丰富的语义信息。
• 高效架构：基于预训练的MiniLM-L12模型进行微调，结合高效的对比学习目标，在保证性能的同时提高训练效率。
• 多领域适用性：训练数据涵盖了多种不同类型的数据集，包括问答、图像描述、代码搜索等，使模型具有广泛的适用性。

📦 安装指南

使用此模型前，需安装SentenceTransformers库，可通过以下命令进行安装：

pip install sentence-transformers

💻 使用示例

基础用法

以下是如何使用SentenceTransformers库来获取给定文本特征的示例代码：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v3_MiniLM-L12')
text = "Replace me by any text you'd like."
text_embbedding = model.encode(text)
# array([-0.01559514,  0.04046123,  0.1317083 ,  0.00085931,  0.04585106,
#        -0.05607086,  0.0138078 ,  0.03569756,  0.01420381,  0.04266302 ...],
#        dtype=float32)

📚 详细文档

训练过程

预训练

我们使用了预训练的MiniLM-L12模型。有关预训练过程的更多详细信息，请参考该模型的卡片。

微调

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

超参数

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

训练数据

我们使用多个数据集的组合来微调模型。句子对的总数超过10亿。每个数据集的采样概率在data_config.json文件中有详细配置。

🔧 技术细节

• 模型架构：基于MiniLM-L12预训练模型，通过微调适应句子嵌入任务。
• 对比学习目标：计算句子对的余弦相似度，并使用交叉熵损失进行优化，使相似句子的向量在嵌入空间中更接近。
• 硬件加速：使用7个TPU v3 - 8进行训练，并得到了Google的Flax、JAX和云团队成员在高效深度学习框架方面的支持。