BERTA-onnx

ONNX-сборка sergeyzh/BERTA — русскоязычной модели эмбеддингов (128M параметров, 768 измерений, до 512 токенов).

Веса не менялись: это тот же экспорт optimum в fp32, чтобы модель можно было запускать одним onnxruntime, без PyTorch. Пригодится там, где нужен семантический поиск по русскому тексту, а тянуть torch не хочется — например, в локальных инструментах и утилитах.

Что внутри

Файл	Зачем
model.onnx	сама модель, fp32, 490 МБ
tokenizer.json	токенизатор
embed.json	параметры для рантайма: пулинг, префиксы, размерность
config.json, tokenizer_config.json, special_tokens_map.json	как в оригинале

Пулинг — mean по маске внимания, затем L2-нормализация. Префиксы у модели свои: search_query: для запроса и search_document: для документа.

Как пользоваться

import json
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

cfg = json.load(open("embed.json", encoding="utf-8"))
tok = Tokenizer.from_file("tokenizer.json")
tok.enable_truncation(max_length=cfg["max_seq_length"])
tok.enable_padding()
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])


def encode(texts, kind="document"):
    enc = tok.encode_batch([cfg["prefixes"][kind] + t for t in texts])
    ids = np.array([e.ids for e in enc], dtype=np.int64)
    mask = np.array([e.attention_mask for e in enc], dtype=np.int64)
    feed = {"input_ids": ids, "attention_mask": mask}
    names = {i.name for i in sess.get_inputs()}
    if "token_type_ids" in names:
        feed["token_type_ids"] = np.zeros_like(ids)
    out = sess.run(None, {k: v for k, v in feed.items() if k in names})[0]
    m = mask[..., None].astype(np.float32)
    vec = (out * m).sum(1) / np.clip(m.sum(1), 1e-9, None)
    return vec / np.linalg.norm(vec, axis=1, keepdims=True)


docs = encode(["Кириллица не рисуется в билде: шрифт без глифов"])
query = encode(["почему не отображаются русские буквы"], kind="query")
print(float(query[0] @ docs[0]))

Проверка качества

Сборка проверялась на задаче поиска по русским техническим заметкам (карточки задач, комментарии разработчиков). Метрика считалась на связях, которые уже были в данных: заголовок задачи ищет своё описание, комментарий ищет задачу, к которой написан.

	recall@1	recall@5
подстрочный поиск	52.9%	75.2%
paraphrase-multilingual-MiniLM-L12-v2	42.6%	68.1%
rubert-tiny-turbo	53.2%	77.1%
LaBSE-ru-turbo	70.6%	90.6%
BERTA (эта модель)	77.4%	90.6%

Мультиязычные модели на таком материале проигрывали даже обычному поиску по подстроке — отсюда и выбор русской модели.

Квантование int8 брать не стоит: размер падает с 490 до 125 МБ, но recall@1 падает с 77.4% до 65.2%, а на поиске по комментариям опускается ниже подстрочного baseline.

Экспорт в ONNX качество не меняет: метрики сборки в точности совпали с метриками оригинала через sentence-transformers.

Лицензия

MIT, как у оригинальной модели. Вся заслуга за качество — у авторов sergeyzh/BERTA и LaBSE-ru-turbo, лежащей в её основе.