BERTA

🚀 BERTA

A model for calculating sentence embeddings in Russian and English, obtained through distillation.

This model calculates sentence embeddings in Russian and English. It is created by distilling the embeddings of ai-forever/FRIDA (embedding size - 1536, 24 layers) into sergeyzh/LaBSE-ru-turbo (embedding size - 768, 12 layers). The main usage mode of FRIDA, CLS pooling, is replaced with mean pooling. No other changes to the model's behavior are made. The distillation is carried out to the maximum extent - for Russian and English sentence embeddings and prefix operations.

The model's context size is the same as FRIDA - 512 tokens.

✨ Features

Prefixes

All prefixes are inherited from FRIDA. The optimal prefix (providing average results) for most tasks, "categorize_entailment: ", is set as the default in config_sentence_transformers.json.

The list of used prefixes and their influence on the model's evaluations in encodechka:

Tasks:

• Semantic text similarity (STS);
• Paraphrase identification (PI);
• Natural language inference (NLI);
• Sentiment analysis (SA);
• Toxicity identification (TI).

Metrics

The model's evaluations on the ruMTEB benchmark:

💻 Usage Examples

Usage with the transformers library

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")
model = AutoModel.from_pretrained("sergeyzh/BERTA")

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.9530372023582458, 0.866746723651886,  0.7839133143424988]
# [0.9360030293464661, 0.8591322302818298, 0.728583037853241] - FRIDA

Usage with sentence_transformers (sentence-transformers>=2.4.0)

from sentence_transformers import SentenceTransformer

# loads model with mean pooling
model = SentenceTransformer("sergeyzh/BERTA")

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

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

📄 License

This project is licensed under the MIT license.