Stella Base En V2

🚀 stella模型

stella是一个通用的文本编码模型，可将文本转换为向量表示，用于检索、聚类、分类等自然语言处理任务。该模型主要有以下几种：

完整的训练思路和训练过程已记录在博客1和博客2，欢迎阅读讨论。

✨ 主要特性

• 多语言支持：涵盖中文和英文，适用于不同语言场景。
• 使用便捷：部分模型无需额外的前缀文本，降低使用门槛。
• 长文本处理：具备一定的长文本编码能力，通过特殊处理规则优化效果。
• 持续更新：不断推出新的版本和模型，提升性能和功能。

📦 安装指南

在使用stella模型前，你需要安装相关依赖库。可以使用以下命令进行安装：

pip install sentence-transformers transformers sklearn torch numpy mteb

💻 使用示例

stella 中文系列模型

在sentence-transformer库中的使用方法：

from sentence_transformers import SentenceTransformer

sentences = ["数据1", "数据2"]
model = SentenceTransformer('infgrad/stella-base-zh-v2')
print(model.max_seq_length)
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

直接使用transformers库：

from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

model = AutoModel.from_pretrained('infgrad/stella-base-zh-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh-v2')
sentences = ["数据1", "数据ABCDEFGH"]
batch_data = tokenizer(
    batch_text_or_text_pairs=sentences,
    padding="longest",
    return_tensors="pt",
    max_length=1024,
    truncation=True,
)
attention_mask = batch_data["attention_mask"]
model_output = model(**batch_data)
last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = normalize(vectors, norm="l2", axis=1, )
print(vectors.shape)  # 2,768

stella 英文系列模型

使用Sentence-Transformers：

from sentence_transformers import SentenceTransformer

sentences = ["one car come", "one car go"]
model = SentenceTransformer('infgrad/stella-base-en-v2')
print(model.max_seq_length)
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

使用HuggingFace Transformers：

from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

model = AutoModel.from_pretrained('infgrad/stella-base-en-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-en-v2')
sentences = ["one car come", "one car go"]
batch_data = tokenizer(
    batch_text_or_text_pairs=sentences,
    padding="longest",
    return_tensors="pt",
    max_length=512,
    truncation=True,
)
attention_mask = batch_data["attention_mask"]
model_output = model(**batch_data)
last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = normalize(vectors, norm="l2", axis=1, )
print(vectors.shape)  # 2,768

📚 详细文档

训练数据

1. 开源数据：选用了wudao_base_200GB[1]、m3e[2]和simclue[3]等开源数据，并着重挑选了长度大于512的文本。
2. 构造数据：在通用语料库上使用LLM构造了一批(question, paragraph)和(sentence, paragraph)数据。

训练方法

1. 对比学习损失函数
2. 带有难负例的对比学习损失函数：分别基于bm25和vector构造了难负例。
3. EWC(Elastic Weights Consolidation)[4]
4. cosent loss[5]
5. 迭代更新：每一种类型的数据一个迭代器，分别计算loss进行更新。

初始权重

stella-base-zh和stella-large-zh分别以piccolo-base-zh[6]和piccolo-large-zh作为基础模型，512 - 1024的position embedding使用层次分解位置编码[7]进行初始化。感谢商汤科技研究院开源的piccolo系列模型。

stella-v2改进

stella-v2在stella模型的基础上，使用了更多的训练数据，同时通过知识蒸馏等方法去除了前置的instruction（比如piccolo的查询:, 结果:, e5的query:和passage:）。

🔧 技术细节

硬件

单卡A100 - 80GB

环境

torch1.13.*; transformers-trainer + deepspeed + gradient-checkpointing

学习率

1e - 6

batch_size

• base模型为1024，额外增加20%的难负例。
• large模型为768，额外增加20%的难负例。

数据量

• 第一版模型约100万，其中用LLM构造的数据约有200K。LLM模型大小为13b。
• v2系列模型到了2000万训练数据。

📄 许可证

本项目采用MIT许可证。

📈 评测指标

C-MTEB leaderboard (Chinese)

MTEB leaderboard (English)

复现结果

C-MTEB:

import torch
import numpy as np
from typing import List
from mteb import MTEB
from sentence_transformers import SentenceTransformer


class FastTextEncoder():
    def __init__(self, model_name):
        self.model = SentenceTransformer(model_name).cuda().half().eval()
        self.model.max_seq_length = 512

    def encode(
            self,
            input_texts: List[str],
            *args,
            **kwargs
    ):
        new_sens = list(set(input_texts))
        new_sens.sort(key=lambda x: len(x), reverse=True)
        vecs = self.model.encode(
            new_sens, normalize_embeddings=True, convert_to_numpy=True, batch_size=256
        ).astype(np.float32)
        sen2arrid = {sen: idx for idx, sen in enumerate(new_sens)}
        vecs = vecs[[sen2arrid[sen] for sen in input_texts]]
        torch.cuda.empty_cache()
        return vecs


if __name__ == '__main__':
    model_name = "infgrad/stella-base-zh-v2"
    output_folder = "zh_mteb_results/stella-base-zh-v2"
    task_names = [t.description["name"] for t in MTEB(task_langs=['zh', 'zh-CN']).tasks]
    model = FastTextEncoder(model_name)
    for task in task_names:
        MTEB(tasks=[task], task_langs=['zh', 'zh-CN']).run(model, output_folder=output_folder)

MTEB: 你可以使用官方脚本复现结果。scripts/run_mteb_english.py

长文本评测

现有数据集在评估模型长文本编码能力方面存在问题，如长度大于512的文本过少，且多数情况下检索只需前512的文本内容。为解决此问题，搜集整理了6份长文本测试集：

• CMRC2018，通用百科
• CAIL，法律阅读理解
• DRCD，繁体百科，已转简体
• Military，军工问答
• Squad，英文阅读理解，已转中文
• Multifieldqa_zh，清华的大模型长文本理解能力评测数据[9]

处理规则是选取答案在512长度之后的文本，短的测试数据会欠采样，长短文本占比约为1:2。除Military数据集外，其他5个测试数据的下载地址：https://drive.google.com/file/d/1WC6EWaCbVgz-vPMDFH4TwAMkLyh5WNcN/view?usp=sharing

评测指标为Recall@5，结果如下：

📋 ToDoList

• 评测的稳定性：评测过程中Clustering任务与官方结果有±0.0x的小差距，因聚类代码未设置random_seed，差距可忽略。
• 更高质量的长文本训练和测试数据：训练数据存在噪声，测试数据问题类型单一，不符合真实分布。
• OOD的性能：在非通用领域，包括stella在内的众多模型效果不如BM25。

📖 参考资料

1. https://www.scidb.cn/en/detail?dataSetId=c6a3fe684227415a9db8e21bac4a15ab
2. https://github.com/wangyuxinwhy/uniem
3. https://github.com/CLUEbenchmark/SimCLUE
4. https://arxiv.org/abs/1612.00796
5. https://kexue.fm/archives/8847
6. https://huggingface.co/sensenova/piccolo-base-zh
7. https://kexue.fm/archives/7947
8. https://github.com/FlagOpen/FlagEmbedding
9. https://github.com/THUDM/LongBench

🌟 新闻动态

• [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

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

⚠️ 重要提示

因为长文本评测数据数量稀少，所以构造时也使用了train部分，如果自行评测，请注意模型的训练数据以免数据泄露。