Piccolo Base Zh

🚀 piccolo-base-zh

piccolo 是一个中文通用文本嵌入模型，由商汤科技研究院通用模型组训练完成。该模型借鉴了 E5 和 GTE 的训练流程，采用两阶段训练方式，目前提供了 piccolo-base-zh 和 piccolo-large-zh 两个版本。

🚀 快速开始

在 sentence-transformer package 中可以很容易地调用 piccolo 模型。

# for s2s dataset, you can use piccolo as below
# 对于短对短数据集，下面是通用的使用方式
from sentence_transformers import SentenceTransformer
sentences = ["数据1", "数据2"]
model = SentenceTransformer('sensenova/piccolo-base-zh')
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

# for s2p dataset, we recommend to add instruction for passage retrieval
# 对于短对长数据集，我们推荐添加 instruction，来帮助模型更好地进行检索。
from sentence_transformers import SentenceTransformer
queries = ['query_1', 'query_2']
passages = ["doc_1", "doc_2"]

model = SentenceTransformer('sensenova/piccolo-base-zh')
q_embeddings = model.encode(["查询：" + q for q in queries], normalize_embeddings=True)
p_embeddings = model.encode(["结果：" + p for p in passages], normalize_embeddings=True)
scores = q_embeddings @ p_embeddings.T

✨ 主要特性

• 两阶段训练：第一阶段搜集和爬取 4 亿中文文本对（弱监督文本对数据），采用二元组的 softmax 对比学习损失优化模型；第二阶段搜集整理 2000 万人工标注的中文文本对（精标数据），采用带有难负样本的三元组的 softmax 对比学习损失进一步优化。
• 多场景适用：在分类、聚类、重排序、检索和语义文本相似度等多种任务上均有良好表现。

📦 安装指南

文档未提及安装步骤，故跳过此章节。

💻 使用示例

基础用法

# for s2s dataset, you can use piccolo as below
# 对于短对短数据集，下面是通用的使用方式
from sentence_transformers import SentenceTransformer
sentences = ["数据1", "数据2"]
model = SentenceTransformer('sensenova/piccolo-base-zh')
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

高级用法

# for s2p dataset, we recommend to add instruction for passage retrieval
# 对于短对长数据集，我们推荐添加 instruction，来帮助模型更好地进行检索。
from sentence_transformers import SentenceTransformer
queries = ['query_1', 'query_2']
passages = ["doc_1", "doc_2"]

model = SentenceTransformer('sensenova/piccolo-base-zh')
q_embeddings = model.encode(["查询：" + q for q in queries], normalize_embeddings=True)
p_embeddings = model.encode(["结果：" + p for p in passages], normalize_embeddings=True)
scores = q_embeddings @ p_embeddings.T

📚 详细文档

模型对比

我们将 piccolo 与其他的开源 embedding 模型在 CMTEB 榜单上进行了比较，请参考 CMTEB 榜单。我们在 eval 文件夹中提供了复现结果的脚本。

训练细节

预训练

pretrain 通常不需要太大的 max length，推荐 128。小的 max length 用以提高 batch size，加快训练速度，从而适应大规模数据。pretrain 损失我们采用二元组 contrastive loss，不加入 hard negative，直接采用 inbatch negative。在实际训练中，我们使用了 32 张 40G A100 进行训练，单卡的 batch size 为 1024。

微调

对于微调，我们通常将 max length 扩增到 512，以适应更大长度的文本输入。微调时会更多地采样 S2P 数据，以增强模型在检索任务上的性能。微调损失采用三元组 contrastive loss，加入 hard negative，neg num 通常设置为 2 - 7，损失计算方式可以参考 GTE 里的 improved contrastive loss。注意，我们给 query 和 passage 设置了不同的 max length，query 的 max length 始终保持在 64。

其他技巧

1. 减小显存的方式：fp16 + gradient checkpointing + ZERO STAGE1（stage2 不支持双塔结构下的 gradient checkpointing），相关 issue 见：https://github.com/microsoft/DeepSpeed/issues/988
2. 数据集采样器：我们采用了 M3E 的 dataset sampler，用以保证每个 batch 里的样本均来自于一个 dataset，负样本更有价值。
3. 指令：指令在我们的实验中对检索任务有非常大的性能提升，我们在每个训练样本前都加入'查询: '和'结果: '这样的指令。

🔧 技术细节

文档中关于模型训练的技术细节描述较为详细，包括训练阶段、损失函数、显存优化技巧等，满足技术细节章节的要求。

📄 许可证

Piccolo 使用 MIT License，可免费用于商业目的。

致谢

piccolo 由商汤科技研究院通用模型组训练完成。Jinkin 完成了代码实现和模型训练，Jinkin 和 CCCCxxx 共同完成了数据搜集、整理和评测工作。项目由 Gaomengya 和 chaorenwu111 主导。同时，感谢 lux0933 以及 yangkai001 的交流与讨论，他们提供了许多有用的建议。