DaniShe/BERTA-onnx

BERTA-onnx

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

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

Что внутри

Пулинг — 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]))

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

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

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

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

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

Лицензия

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