Dunzhang Stella En 400M V5

🚀 Marqo Stella v2

Marqo Stella v2 是一个与原始 Dunzhang stella 400m 模型 相似的模型，它融合了一个俄罗斯套娃层（Matryoshka Layer）。这种层级结构能够在不改变相关性指标的前提下，降低生成嵌入向量时的计算开销。

🚀 快速开始

环境准备

确保你已经安装了必要的库，如 transformers、torch 和 sklearn。可以使用以下命令进行安装：

pip install transformers torch sklearn

代码示例

以下是一个使用该模型进行查询和文档嵌入，并计算相似度的示例代码：

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

# 定义查询提示
query_prompt = "Instruct: Given a web search query, retrieve relevant passages that answer the query.\nQuery: "
# 定义查询列表
queries = [
    "What are some ways to reduce stress?",
    "What are the benefits of drinking green tea?",
]
# 为每个查询添加提示
queries = [query_prompt + query for query in queries]
# 定义文档列表，文档不需要提示
docs = [
    "There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent stress from building up.",
    "Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.",
]

# 克隆模型后的本地路径
model_dir = "Marqo/dunzhang-stella_en_400M_v5"
# 加载模型并将其移动到 GPU 上进行评估
model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).cuda().eval()
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

# 对查询进行嵌入
with torch.no_grad():
    input_data = tokenizer(queries, padding="longest", truncation=True, max_length=512, return_tensors="pt")
    input_data = {k: v.cuda() for k, v in input_data.items()}
    attention_mask = input_data["attention_mask"]
    last_hidden_state = model(**input_data)[0]
    last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
    query_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    query_vectors = normalize(query_vectors.cpu().numpy())

# 对文档进行嵌入
with torch.no_grad():
    input_data = tokenizer(docs, padding="longest", truncation=True, max_length=512, return_tensors="pt")
    input_data = {k: v.cuda() for k, v in input_data.items()}
    attention_mask = input_data["attention_mask"]
    last_hidden_state = model(**input_data)[0]
    last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
    docs_vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
    docs_vectors = normalize(docs_vectors.cpu().numpy())

# 打印查询向量和文档向量的形状
print(query_vectors.shape, docs_vectors.shape)
# (2, 1024) (2, 1024)

# 计算查询向量和文档向量之间的相似度
similarities = query_vectors @ docs_vectors.T
print(similarities)
# [[0.8397531  0.29900077]
#  [0.32818374 0.80954516]]

💻 使用示例

基础用法

上述代码展示了如何使用该模型进行查询和文档的嵌入，并计算它们之间的相似度。具体步骤如下：

1. 定义查询和文档：准备需要处理的查询和文档列表。
2. 加载模型和分词器：从指定路径加载模型和分词器。
3. 对查询和文档进行嵌入：使用分词器对查询和文档进行分词，并通过模型生成嵌入向量。
4. 计算相似度：使用矩阵乘法计算查询向量和文档向量之间的相似度。

高级用法

你可以根据实际需求对代码进行扩展，例如：

• 批量处理：处理更多的查询和文档，提高效率。
• 不同的相似度计算方法：除了矩阵乘法，还可以使用其他相似度计算方法，如余弦相似度。
• 与其他模型结合使用：将该模型的输出与其他模型的输出进行融合，以获得更好的性能。

📄 许可证

本项目采用 MIT 许可证。

模型评估结果

以下是该模型在多个数据集上的评估结果：