Модель для расчетов эмбеддингов предложений на русском и английском языках получена методом дистилляции эмбеддингов
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:
Оценки модели на бенчмарке
ruMTEB:
1import torch
2import torch.nn.functional as F
3from transformers import AutoTokenizer, AutoModel
4
5
6def pool(hidden_state, mask, pooling_method="mean"):
7 if pooling_method == "mean":
8 s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
9 d = mask.sum(axis=1, keepdim=True).float()
10 return s / d
11 elif pooling_method == "cls":
12 return hidden_state[:, 0]
13
14inputs = [
15 #
16 "paraphrase: В Ярославской области разрешили работу бань, но без посетителей",
17 "categorize_entailment: Женщину доставили в больницу, за ее жизнь сейчас борются врачи.",
18 "search_query: Сколько программистов нужно, чтобы вкрутить лампочку?",
19 #
20 "paraphrase: Ярославским баням разрешили работать без посетителей",
21 "categorize_entailment: Женщину спасают врачи.",
22 "search_document: Чтобы вкрутить лампочку, требуется три программиста: один напишет программу извлечения лампочки, другой — вкручивания лампочки, а третий проведет тестирование."
23]
24
25tokenizer = AutoTokenizer.from_pretrained("sergeyzh/rubert-mini-frida")
26model = AutoModel.from_pretrained("sergeyzh/rubert-mini-frida")
27
28tokenized_inputs = tokenizer(inputs, max_length=512, padding=True, truncation=True, return_tensors="pt")
29
30with torch.no_grad():
31 outputs = model(**tokenized_inputs)
32
33embeddings = pool(
34 outputs.last_hidden_state,
35 tokenized_inputs["attention_mask"],
36 pooling_method="mean"
37)
38
39embeddings = F.normalize(embeddings, p=2, dim=1)
40sim_scores = embeddings[:3] @ embeddings[3:].T
41print(sim_scores.diag().tolist())
42# [0.9423348903656006, 0.8306248188018799, 0.7095720767974854]
43# [0.9360030293464661, 0.8591322302818298, 0.728583037853241] - FRIDA
1from sentence_transformers import SentenceTransformer
2
3# loads model with mean pooling
4model = SentenceTransformer("sergeyzh/rubert-mini-frida")
5
6paraphrase = model.encode(["В Ярославской области разрешили работу бань, но без посетителей", "Ярославским баням разрешили работать без посетителей"], prompt="paraphrase: ")
7print(paraphrase[0] @ paraphrase[1].T)
8# 0.94233495
9# 0.9360032 - FRIDA
10
11categorize_entailment = model.encode(["Женщину доставили в больницу, за ее жизнь сейчас борются врачи.", "Женщину спасают врачи."], prompt="categorize_entailment: ")
12print(categorize_entailment[0] @ categorize_entailment[1].T)
13# 0.8306249
14# 0.8591322 - FRIDA
15
16query_embedding = model.encode("Сколько программистов нужно, чтобы вкрутить лампочку?", prompt="search_query: ")
17document_embedding = model.encode("Чтобы вкрутить лампочку, требуется три программиста: один напишет программу извлечения лампочки, другой — вкручивания лампочки, а третий проведет тестирование.", prompt="search_document: ")
18print(query_embedding @ document_embedding.T)
19# 0.70957196
20# 0.7285831 - FRIDA
