ru-reranker-edge-150m

Русский кросс-энкодер-реранкер на 150M параметров. Обучен дистилляцией из BAAI/bge-reranker-v2-m3 (568M) и на русских задачах даёт 95–99% его качества при в 3.8 раза меньшем размере.

Что это

• База: deepvk/RuModernBERT-base (контекст до 8192 токенов, обучали и оценивали при 512).
• Как обучали: листвайз-дистилляция — модель учится повторять распределение скоров учителя внутри группы «запрос + 16 документов». Меток релевантности не требуется.
• На чём обучали: русский mMARCO (train), MIRACL-ru (train), статьи русской Википедии, фрагменты кодексов и нормативных актов РФ, синтетические пары «аргумент → контраргумент» (Qwen3-27B). Ни один из оценочных датасетов ниже в обучение не попадал.
• Вход: пара (запрос, документ) → один логит. Чем больше, тем релевантнее. sentence-transformers и vLLM по умолчанию пропускают логит через сигмоиду (скор 0…1); transformers, ONNX, Transformers.js и TEI с raw_scores: true отдают сырой логит.

Качество

nDCG@10 на подмножестве RusBEIR (зеро-шот, реранк BM25 top-100, одинаковые кандидаты во всех строках):

RuBQReranking (ruMTEB, официальный mteb, зеро-шот): MAP 0.773, nDCG@10 0.839 (учитель: 0.785 / 0.848).

Среднее по пяти датасетам — выше учителя; ArguAna (поиск контраргументов) закрыт синтетическими парами «аргумент → контраргумент» (0.249 → 0.386).

Режимы работы

Все варианты ниже проверены на одних и тех же парах; расхождение — максимум |Δ логита| относительно transformers fp32.

Работает и с transformers 4.x, и с 5.x.

sentence-transformers

from sentence_transformers import CrossEncoder

model = CrossEncoder("sshalimov04/ru-reranker-edge-150m", max_length=512)
query = "Что может вызвать цунами?"
docs = [
    "Землетрясения, извержения вулканов и подводные оползни обладают потенциалом вызвать цунами.",
    "Флаг муниципального округа утверждён решением совета депутатов.",
]
scores = model.predict([(query, d) for d in docs])   # сигмоида от логита, 0…1
ranked = model.rank(query, docs)                      # [{'corpus_id': 0, 'score': ...}, ...]

Сырые логиты: CrossEncoder(..., activation_fn=torch.nn.Identity()).

На CPU без torch — ONNX или OpenVINO (pip install "sentence-transformers[onnx]" / [openvino]):

model = CrossEncoder("sshalimov04/ru-reranker-edge-150m", backend="onnx")  # onnx/model.onnx
model = CrossEncoder("sshalimov04/ru-reranker-edge-150m", backend="onnx",
                     model_kwargs={"file_name": "onnx/model_quantized.onnx"})  # int8, 144 МБ
model = CrossEncoder("sshalimov04/ru-reranker-edge-150m", backend="openvino")

transformers

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

name = "sshalimov04/ru-reranker-edge-150m"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name).eval()
enc = tok([query] * len(docs), docs, padding=True, truncation=True, max_length=512, return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits.squeeze(-1)

Или pipeline("text-classification", model=name) с входом {"text": query, "text_pair": doc} и function_to_apply="none".

Text Embeddings Inference

docker run -p 8080:80 ghcr.io/huggingface/text-embeddings-inference:cpu-1.9 \
    --model-id sshalimov04/ru-reranker-edge-150m
# GPU: образ ghcr.io/huggingface/text-embeddings-inference:1.9 (или тег под вашу архитектуру) и --gpus all

curl localhost:8080/rerank -H 'Content-Type: application/json' \
    -d '{"query": "Что может вызвать цунами?", "texts": ["...", "..."], "raw_scores": true}'

CPU-образ берёт onnx/model.onnx, GPU-образ — model.safetensors (Candle). Candle считает GELU через tanh-аппроксимацию, отсюда расхождение ~0.002; на порядок документов это не влияет.

vLLM

vllm serve sshalimov04/ru-reranker-edge-150m --runner pooling
curl localhost:8000/rerank -H 'Content-Type: application/json' \
    -d '{"query": "Что может вызвать цунами?", "documents": ["...", "..."]}'

from vllm import LLM
llm = LLM("sshalimov04/ru-reranker-edge-150m", runner="pooling")
scores = [o.outputs.score for o in llm.score(query, docs)]   # сигмоида от логита

Проверено на vLLM 0.33 (CPU-образ vllm/vllm-openai-cpu).

Transformers.js (браузер / Node)

import { AutoTokenizer, AutoModelForSequenceClassification } from "@huggingface/transformers";

const name = "sshalimov04/ru-reranker-edge-150m";
const tok = await AutoTokenizer.from_pretrained(name);
const model = await AutoModelForSequenceClassification.from_pretrained(name, { dtype: "q8" }); // "fp32" | "fp16" | "q8"
const inputs = tok(new Array(docs.length).fill(query), { text_pair: docs, padding: true, truncation: true, max_length: 512 });
const { logits } = await model(inputs);

Рабочий пример — демо в браузере.

ONNX Runtime напрямую

Вход: input_ids, attention_mask (int64, [batch, seq]); выход: logits [batch, 1]. Оси batch и seq динамические, проверено на длинах от 16 до 512 токенов.

Скорость

На NVIDIA GB10 (bf16, batch 32, длина 512) реранк 100 документов: ~0.8 с в обычном transformers-пути и ~0.35–0.4 с с torch.compile (2.2× быстрее, без потерь):

model.model = torch.compile(model.model, dynamic=True)   # для CrossEncoder

ONNX и квантизация

• onnx/model.onnx — fp32, opset 17, совпадает с torch до 1e-4.
• onnx/model_fp16.onnx — fp16, Spearman 0.997 по скорам. На CUDA через ONNX Runtime медленнее torch (часть узлов ModernBERT уходит на CPU), поэтому он нужен для переноса на другие рантаймы и устройства, а не для ускорения на GPU.
• onnx/model_quantized.onnx — int8, поканальная динамическая квантизация (per_channel=True, reduce_range=True, MatMul и таблица эмбеддингов), 144 МБ вместо 571. Ранжирование не портит: на отложенных 500 юридических делах gold@5 0.720 против 0.702 у fp32 при длине 512. Обычная (не поканальная) динамическая int8 ломает скоры: Spearman падает до 0.77–0.95, поэтому она не выложена.

Ограничения

• Только русский язык (тексты на других языках не тестировались).
• Путь внимания ModernBERT в transformers без FlashAttention — узкое место; torch.compile его частично закрывает.
• На юридическом корпусе с номерами статей один реранкер BM25 не обыгрывает. На отложенных 500 делах gold@5: BM25 0.736, реранкер 0.702. Запросы содержат номера статей, точное совпадение токенов сильнее семантики. Рабочий рецепт — смесь нормированных скоров 50/50: gold@5 0.760, gold@10 0.806 против 0.774 у BM25 (парный тест Макнемара p=0.0015). Одну модель как единственную первую стадию в таких доменах ставить не стоит.
• Оценка первой стадии — bm25s, а не Elasticsearch, поэтому абсолютные числа могут отличаться от лидерборда RusBEIR; сравнения внутри таблицы точны.

Данные для воспроизведения

Кэш скоров учителя и тексты всех обучающих групп опубликованы как датасет: sshalimov04/ru-reranker-teacher-scores.

---

English. A 150M Russian cross-encoder reranker (RuModernBERT-base) distilled listwise from BAAI/bge-reranker-v2-m3. Zero-shot on a RusBEIR subset it averages 0.609 nDCG@10 (BM25 0.517, teacher 0.607) and reaches MAP 0.773 on RuBQReranking (teacher 0.785). Input is a (query, document) pair, output a relevance logit (sentence-transformers applies a sigmoid by default). Runs with sentence-transformers (torch / ONNX / OpenVINO), transformers 4.x and 5.x, Text Embeddings Inference, vLLM, Transformers.js (fp32 / fp16 / q8) and ONNX Runtime. Teacher scores and training groups are released as sshalimov04/ru-reranker-teacher-scores.