新闻 | News

[2024-04-06] 开源 puff 系列模型，专为检索与语义匹配任务优化，侧重泛化能力与私有通用测试集表现，支持可变向量维度，兼容中英双语。

[2024-02-27] 开源 stella-mrl-large-zh-v3.5-1792d 模型，支持向量维度动态调整。

[2024-02-17] 开源 stella v3 系列、对话专用编码模型及配套训练数据。

[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

欢迎访问本人主页获取最新模型，并留下宝贵建议！

1 开源资源列表

本次开源包含两个通用向量编码模型和一个对话专用向量编码模型，同时开放全量160万对话重写数据集及20万难负例检索数据集。

开源模型：

开源数据：

1. 全量对话重写数据集（约160万条）
2. 带难负例的检索数据集（约20万条）

以上数据均采用LLM技术构建，诚邀社区共同贡献数据资源。

2 使用指南

2.1 通用编码模型使用方法

通过SentenceTransformer直接加载：

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"])

2.2 对话编码模型使用方法

使用场景： 在一段对话中，需要根据用户语句检索相关文本时，若对话存在大量指代和省略现象导致通用编码模型效果不佳，可采用本项目专用的对话编码模型进行语义编码

使用要点：

1. 对话编码时，每个话语需遵循格式："{角色}: {文本内容}"，并通过[SEP]符号进行拼接
2. 需将完整对话输入模型编码，若长度超限则优先删除早期对话记录。需特别注意：最终生成的向量本质是对末句对话的重写版本编码结果！
3. 对话内容使用stella-dialogue-large-zh-v3-1792d编码，被检索文本使用stella-large-zh-v3-1792d编码，因此本场景需要同时使用两个编码模型

若对使用方法存在疑问，请参阅后续章节了解该模型的训练机制。

使用示例：

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)

3 通用编码模型训练技巧分享

Hard Negative

难负例挖掘也是一个经典的训练技巧，几乎总能有效提升模型性能

Dropout-1D

Dropout早已成为深度学习中的标准配置，我们可以稍作调整使其更适用于句向量训练。 在训练过程中，我们尝试让每个词嵌入都能表征整个句子的语义，而在推理阶段通过均值池化实现类似模型集成的效果。 具体实现是在均值池化操作中加入一维dropout，PyTorch代码如下：

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]

4 对话编码模型细节

4.1 为什么需要对话编码模型？

参见本人历史文章：https://www.zhihu.com/pin/1674913544847077376

4.2 训练数据

单条数据示例：

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

4.3 训练损失

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

dial_model 是指待训练的模型，本人以 stella-large-zh-v3-1792d 作为基础模型（base-model）进行继续训练。

existing_model 指现已训练完成的通用编码模型，本人使用的是 stella-large-zh-v3-1792d。

目前已开源 dialogue-embedding 的全量训练数据，理论上可复现本模型效果。

Loss 下降情况如下：

4.4 效果

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

5 后续计划

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

6 常见问题

Q: 为什么向量维度是1792？
A: 最初计划发布768、1024、768+768、1024+1024、1024+768等多种维度，但因时间有限，目前仅完成1792维度的模型。理论上维度越高效果越好。

Q: 如何复现 CMTEB 评测效果？
A: 通过 SentenceTransformer 加载模型后直接使用官方评测脚本即可。请注意，分类任务中的向量需先进行归一化（normalize）处理。

Q: 复现的 CMTEB 结果与本文报告不一致？
A: 聚类结果不一致属正常现象，因官方评测代码未设定随机种子。若其他项目存在差异，建议检查代码或联系本人。

Q: 如何选择向量模型？
A: 没有一劳永逸的解决方案，建议在自己的测试集上验证效果。本人推荐 bge、e5 和 stella 系列模型。

Q: 文本长度为何限定512？能否扩展？
A: 可以扩展但不推荐。更长的文本普遍效果不佳，这是当前训练方法与数据特性所致，几乎无法根本解决。建议对长文本采用分块处理。

Q: 训练资源与算力需求？
A: 需处理亿级别数据量，单张 A100 显卡需持续训练一个月以上。

改写说明：

• 术语与表达专业化：对技术术语和行业表述进行了统一和规范，确保符合NLP和模型训练领域的习惯。
• 结构及格式严格保持：完整保留原文的章节、列表、代码块、链接等Markdown结构和排版，未做删减或合并。
• 技术细节与变量名原样呈现：所有模型名称、维度、文件名、评测名称等专有名词和变量均未翻译，保持原貌。

如果您有其他风格或特定平台发布的需求，我可以进一步为您调整表达。