Stella Large Zh V2

🚀 stella模型

stella是一个通用的文本编码模型，可将文本转换为向量表示，用于检索、语义匹配等自然语言处理任务。它提供了多种不同版本的模型，以满足不同的应用场景和需求。

🚀 快速开始

安装依赖库

使用stella模型前，需要安装sentence-transformers和transformers库，可使用以下命令进行安装：

pip install sentence-transformers transformers

使用示例

以下是使用sentence-transformers库调用stella-base-zh-v2模型的示例代码：

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-base-en-v2、stella-base-zh-v2和stella-large-zh-v2）不需要任何前缀文本。
• 向量维度可变：部分模型支持向量可变维度，如stella-mrl-large-zh-v3.5-1792d。
• 效果优异：在多个基准测试中表现出色，如C-MTEB和MTEB排行榜。

📦 安装指南

使用pip安装所需的依赖库：

pip install sentence-transformers transformers numpy torch

💻 使用示例

基础用法

stella 中文系列模型

• stella-base-zh 和 stella-large-zh：用法和piccolo完全一致，即在检索重排任务上给query和passage加上查询: 和结果: 。对于短短匹配不需要做任何操作。
• stella-base-zh-v2 和 stella-large-zh-v2：任何使用场景中都不需要加前缀文本。

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 英文系列模型

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、m3e和simclue），着重挑选了长度大于512的文本。
2. 在通用语料库上使用大语言模型（LLM）构造的一批(question, paragraph)和(sentence, paragraph)数据。

损失函数

1. 对比学习损失函数
2. 带有难负例的对比学习损失函数（基于bm25和向量构造难负例）
3. EWC（弹性权重整合）
4. cosent损失

模型权重初始化

stella-base-zh和stella-large-zh分别以piccolo-base-zh和piccolo-large-zh作为基础模型，512 - 1024的位置嵌入使用层次分解位置编码进行初始化。

训练策略

每种类型的数据使用一个迭代器，分别计算损失。

模型改进

基于stella模型，stella-v2使用了更多的训练数据，并通过知识蒸馏去除了指令。

🔧 技术细节

训练环境

• 硬件：单卡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万训练数据。

评测指标

C-MTEB排行榜（中文）

MTEB排行榜（英文）

长文本评测

现有数据集存在两个问题：

1. 长度大于512的文本过少。
2. 即便文本长度大于512，对于检索而言也只需要前512的文本内容，导致无法准确评估模型的长文本编码能力。

为了解决这个问题，搜集了相关开源数据并使用规则进行过滤，最终整理了6份长文本测试集：

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

处理规则是选取答案在512长度之后的文本，短的测试数据会欠采样，长短文本占比约为1:2，所以模型既得理解短文本也得理解长文本。除了Military数据集，提供了其他5个测试数据的下载地址：点击下载

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

待解决问题

• 评测的稳定性：评测过程中发现Clustering任务会和官方的结果不一致，大约有±0.0x的小差距，原因是聚类代码没有设置random_seed，差距可以忽略不计，不影响评测结论。
• 更高质量的长文本训练和测试数据：训练数据多是用13B模型构造的，可能存在噪声。测试数据基本都是从阅读理解（MRC）数据整理来的，问题都是事实性（factoid）类型，不符合真实分布。
• OOD的性能：虽然近期出现了很多向量编码模型，但是对于不是那么通用的领域，包括stella、openai和cohere等模型的效果均比不上BM25。

📄 许可证

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

新闻资讯

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