🚀 stella-large-zh-v3-1792d模型介绍

本项目开源了多个通用向量编码模型和一个针对dialogue进行编码的向量模型，同时提供了相关训练数据。这些模型在检索、语义匹配等任务上表现出色，具有泛化性强、中英双语支持等特点。

🚀 快速开始

模型信息

开源模型清单

开源数据清单

1. 全量对话重写数据集 约160万
2. 部分带有难负例的检索数据集 约20万

模型使用示例

💻 使用示例

基础用法

# 通用编码模型使用方法
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("infgrad/stella-base-zh-v3-1792d")
# model = SentenceTransformer("infgrad/stella-large-zh-v3-1792d")
vectors = model.encode(["text1", "text2"])

高级用法

# dialogue编码模型使用方法
from sentence_transformers import SentenceTransformer

dial_model = SentenceTransformer("infgrad/stella-dialogue-large-zh-v3-1792d")
general_model = SentenceTransformer("infgrad/stella-large-zh-v3-1792d")
# dialogue = ["张三: 吃饭吗", "李四: 等会去"]
dialogue = ["A: 最近去打篮球了吗", "B: 没有"]
corpus = ["B没打篮球是因为受伤了。", "B没有打乒乓球"]
last_utterance_vector = dial_model.encode(["[SEP]".join(dialogue)], normalize_embeddings=True)
corpus_vectors = general_model.encode(corpus, normalize_embeddings=True)
# 计算相似度
sims = (last_utterance_vector * corpus_vectors).sum(axis=1)
print(sims)

✨ 主要特性

• 多模型支持：提供通用向量编码模型和dialogue编码模型，满足不同场景需求。
• 数据丰富：开源大量训练数据，包括对话重写数据集和难负例检索数据集。
• 效果出色：在多个任务和数据集上取得了较好的成绩。

📚 详细文档

通用编码模型训练技巧

hard negative

难负例挖掘是经典技巧，能提升模型效果。

dropout-1d

对dropout进行改造，使其更适合句向量训练。在mean_pooling时加入dropout_1d，代码如下：

vector_dropout = nn.Dropout1d(0.3)  # 算力有限，试了0.3和0.5 两个参数，其中0.3更优
last_hidden_state = bert_model(...)[0]
last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
last_hidden = vector_dropout(last_hidden)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

dialogue编码模型细节

为什么需要dialogue编码模型

在对话场景中，通用编码模型效果不佳，dialogue编码模型可解决此问题。具体可参见本人历史文章。

训练数据

单条数据示例：

{
  "dialogue": [
    "A: 最近去打篮球了吗",
    "B: 没有"
  ],
  "last_utterance_rewrite": "B: 我最近没有去打篮球"
}

训练Loss

loss = cosine_loss( dial_model.encode(dialogue), existing_model.encode(last_utterance_rewrite) )

其中，dial_model 是以 stella-large-zh-v3-1792d 为基础模型继续训练的模型，existing_model 是现有的通用编码模型，这里使用的是 stella-large-zh-v3-1792d。

Loss下降情况：

效果

目前暂无专门测试集，简单测试有效果，部分测试结果见文件 dial_retrieval_test.xlsx。

🔧 技术细节

模型评估结果

以下是模型在多个任务和数据集上的评估结果：

📄 许可证

文档中未提及相关许可证信息。

📋 后续计划

1. 收集更多的dial-rewrite数据。
2. 开发不同EmbeddingDimensions的编码模型。

❓ 常见问题解答

Q: 为什么向量维度是1792？ A: 最初考虑发布768、1024，768+768，1024+1024，1024+768维度，但是时间有限，先做了1792就只发布1792维度的模型。理论上维度越高效果越好。

Q: 如何复现CMTEB效果？ A: SentenceTransformer加载后直接用官方评测脚本就行，注意对于Classification任务向量需要先normalize一下。

Q: 复现的CMTEB效果和本文不一致？ A: 聚类不一致正常，官方评测代码没有设定seed，其他不一致建议检查代码或联系本人。

Q: 如何选择向量模型？ A: 没有免费的午餐，在自己测试集上试试，本人推荐bge、e5和stella。

Q: 长度为什么只有512，能否更长？ A: 可以但没必要，长了效果普遍不好，这是当前训练方法和数据导致的，几乎无解，建议长文本还是走分块。

Q: 训练资源和算力？ A: 亿级别的数据，单卡A100要一个月起步。

📰 新闻动态

• [2024-04-06] 开源puff系列模型，专门针对检索和语义匹配任务，更多的考虑泛化性和私有通用测试集效果，向量维度可变，中英双语。
• [2024-02-27] 开源stella-mrl-large-zh-v3.5-1792d模型，支持向量可变维度。
• [2024-02-17] 开源stella v3系列、dialogue编码模型和相关训练数据。
• [2023-10-19] 开源stella-base-en-v2 使用简单，不需要任何前缀文本。
• [2023-10-12] 开源stella-base-zh-v2和stella-large-zh-v2，效果更好且使用简单，不需要任何前缀文本。
• [2023-09-11] 开源stella-base-zh和stella-large-zh

欢迎去本人主页查看最新模型，并提出您的宝贵意见！