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 許可證。

模型評估結果

以下是該模型在多個數據集上的評估結果：