Stella Base Zh V2

🚀 stella模型

stella是一个通用的文本编码模型，可用于文本编码、检索和语义匹配等任务。它有多种版本，能适应不同的应用场景和需求。

🚀 快速开始

安装依赖

运行以下代码前，请确保你已经安装了必要的库：

pip install torch numpy mteb sentence-transformers transformers sklearn

代码示例

以下是使用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-large-zh-v2、stella-base-zh-v2、stella-large-zh和stella-base-zh等多种模型，以满足不同的需求。
• 无需前缀文本：stella-base-zh-v2和stella-large-zh-v2模型使用简单，在任何使用场景中都不需要加前缀文本。
• 支持向量可变维度：部分模型支持向量可变维度，如stella-mrl-large-zh-v3.5-1792d模型。
• 中英双语支持：部分模型支持中英双语，如puff系列模型。

📦 安装指南

使用sentence-transformers库加载模型：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('infgrad/stella-base-zh-v2')

💻 使用示例

基础用法

from sentence_transformers import SentenceTransformer

sentences = ["这是一个示例句子", "这是另一个示例句子"]
model = SentenceTransformer('infgrad/stella-base-zh-v2')
embeddings = model.encode(sentences)
print(embeddings)

高级用法

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

model = AutoModel.from_pretrained('infgrad/stella-base-zh-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh-v2')
sentences = ["高级用法示例句子", "用于演示更复杂的场景"]
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)

📚 详细文档

模型信息

模型参数

C-MTEB排行榜（中文）

长文本评测

现有数据集存在长度大于512的文本过少，且即便大于512，对于检索而言也只需要前512的文本内容的问题，导致无法准确评估模型的长文本编码能力。为了解决这个问题，整理了6份长文本测试集：

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

除了Military数据集，提供了其他5个测试数据的下载地址：https://drive.google.com/file/d/1WC6EWaCbVgz-vPMDFH4TwAMkLyh5WNcN/view?usp=sharing

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

新闻动态

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

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

🔧 技术细节

训练环境

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

📄 许可证

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

📋 ToDoList

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

⚠️ 重要提示

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

💡 使用建议

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