BERTA-uncased

Модель для расчетов эмбеддингов предложений на русском и английском языках получена методом дистилляции эмбеддингов ai-forever/FRIDA (размер эмбеддингов - 1536, слоёв - 24). Основной режим использования FRIDA - CLS pooling заменен на mean pooling. Дистиляция выполнена с uncased преобразованием текста в максимально возможном объеме - эмбеддинги русских и английских предложений, работа префиксов.

Модель принадлежит к виду uncased - не различает при обработке текста буквы, написанные в верхнем и нижнем регистре (фразы, например, “С Новым Годом!” и “С НОВЫМ ГОДОМ!” кодируются одинаковой последовательностью токенов и имеют равные значения эмбеддингов). Размер эмбеддингов модели - 768, слоёв - 12. Размер контекста модели соответствует FRIDA - 512 токенов.

Префиксы

Все префиксы унаследованы от FRIDA:

• search_query:
• search_document:
• paraphrase:
• categorize:
• categorize_sentiment:
• categorize_topic:
• categorize_entailment:

Метрики

Оценки модели на бенчмарке ruMTEB:

Использование модели с библиотекой 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/BERTA-uncased")
model = AutoModel.from_pretrained("sergeyzh/BERTA-uncased")

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.9447306990623474, 0.8616577386856079, 0.7869422435760498]
# [0.9530372023582458, 0.866746723651886,  0.7839133143424988] - BERTA
# [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/BERTA-uncased")

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

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

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