multilingual-reranker-small, 8-bit

This is Horizon-Labs/multilingual-reranker-small in a smaller ONNX file, for running on a CPU. Same model, same tokenizer, same scores to within a small margin. The file is about a third smaller than the original quantized one.

We use it in Sno Station to rerank memories for coding agents.

What changed

The token embeddings, the attention projections and the first MLP projection are stored in 8 bits. The second MLP projection stays in float. Its inputs carry a few very large values, and 8 bits move the scores too much there.

Weights are unsigned 8-bit, per channel. Signed weights overflow on AVX2 CPUs without VNNI and flatten every score.

quantize.py rebuilds the file from the original onnx/model.onnx at revision 4fd96cc3b19fd47e790cf923879b193d3ebf0bc9. The output is byte-identical.

Use

With transformers.js:

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

const id = "sno-lab/multilingual-reranker-small-int8";
const tokenizer = await AutoTokenizer.from_pretrained(id);
const model = await AutoModelForSequenceClassification.from_pretrained(id, { dtype: "q8" });

const question = "Which database does the project use?";
const memories = ["The project stores everything in PostgreSQL 16.", "We talked about the project schedule on Monday."];
const inputs = tokenizer(memories.map(() => question), { text_pair: memories, padding: true, truncation: true });
const { logits } = await model(inputs);
console.log(logits.data); // higher means more relevant: the first one wins by a wide margin

License

Apache-2.0, the same as the original. The original model is by Horizon Labs and is built on jhu-clsp/mmBERT-small (MIT). See the original model card for its training data and evaluation.