Rubert Mini Frida

🚀 rubert-mini-frida - 輕量快速的FRIDA改進版

rubert-mini-frida是一個用於計算俄語和英語句子嵌入的模型。它通過將ai-forever/FRIDA（嵌入維度為1536，層數為24）的嵌入蒸餾到sergeyzh/rubert-mini-sts（嵌入維度為312，層數為7）中得到。FRIDA的主要使用模式CLS pooling被替換為mean pooling，未對模型行為進行其他更改（如嵌入的修改或過濾、使用額外模型）。蒸餾在最大程度上進行，涵蓋俄語和英語句子的嵌入以及前綴的處理。

模型推薦的上下文大小與FRIDA一致，不超過512個詞元（實際繼承自原始模型的為2048）。

✨ 主要特性

前綴繼承

所有前綴均繼承自FRIDA。對於大多數任務而言，最優（能保證平均效果）的前綴 "categorize: " 已默認寫入 config_sentence_transformers.json。

在 encodechka 中使用的前綴及其對模型評估的影響如下：

任務說明：

• 語義文本相似度（STS）
• 釋義識別（PI）
• 自然語言推理（NLI）
• 情感分析（SA）
• 毒性識別（TI）

模型指標

模型在 ruMTEB 基準測試中的評估結果如下：

📦 安裝指南

文檔未提及安裝步驟，可參考相關庫的官方文檔進行安裝，如 transformers 和 sentence_transformers。

💻 使用示例

基礎用法

使用 transformers 庫

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel


def pool(hidden_state, mask, pooling_method="mean"):
    if pooling_method == "mean":
        s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
        d = mask.sum(axis=1, keepdim=True).float()
        return s / d
    elif pooling_method == "cls":
        return hidden_state[:, 0]

inputs = [
    # 
    "paraphrase: В Ярославской области разрешили работу бань, но без посетителей",
    "categorize_entailment: Женщину доставили в больницу, за ее жизнь сейчас борются врачи.",
    "search_query: Сколько программистов нужно, чтобы вкрутить лампочку?",
    # 
    "paraphrase: Ярославским баням разрешили работать без посетителей",
    "categorize_entailment: Женщину спасают врачи.",
    "search_document: Чтобы вкрутить лампочку, требуется три программиста: один напишет программу извлечения лампочки, другой — вкручивания лампочки, а третий проведет тестирование."
]

tokenizer = AutoTokenizer.from_pretrained("sergeyzh/rubert-mini-frida")
model = AutoModel.from_pretrained("sergeyzh/rubert-mini-frida")

tokenized_inputs = tokenizer(inputs, max_length=512, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():
    outputs = model(**tokenized_inputs)
    
embeddings = pool(
    outputs.last_hidden_state, 
    tokenized_inputs["attention_mask"],
    pooling_method="mean"
)

embeddings = F.normalize(embeddings, p=2, dim=1)
sim_scores = embeddings[:3] @ embeddings[3:].T
print(sim_scores.diag().tolist())
# [0.9423348903656006, 0.8306248188018799, 0.7095720767974854]
# [0.9360030293464661, 0.8591322302818298, 0.728583037853241] - FRIDA

使用 sentence_transformers 庫（sentence-transformers>=2.4.0）

from sentence_transformers import SentenceTransformer

# loads model with mean pooling
model = SentenceTransformer("sergeyzh/rubert-mini-frida")

paraphrase = model.encode(["В Ярославской области разрешили работу бань, но без посетителей", "Ярославским баням разрешили работать без посетителей"], prompt="paraphrase: ")
print(paraphrase[0] @ paraphrase[1].T) 
# 0.94233495
# 0.9360032 - FRIDA

categorize_entailment = model.encode(["Женщину доставили в больницу, за ее жизнь сейчас борются врачи.", "Женщину спасают врачи."], prompt="categorize_entailment: ")
print(categorize_entailment[0] @ categorize_entailment[1].T) 
# 0.8306249
# 0.8591322 - FRIDA

query_embedding = model.encode("Сколько программистов нужно, чтобы вкрутить лампочку?", prompt="search_query: ")
document_embedding = model.encode("Чтобы вкрутить лампочку, требуется три программиста: один напишет программу извлечения лампочки, другой — вкручивания лампочки, а третий проведет тестирование.", prompt="search_document: ")
print(query_embedding @ document_embedding.T) 
# 0.70957196
# 0.7285831 - FRIDA

📄 許可證

本項目採用 MIT 許可證。