mmbert-small-nli-onnx-q8

mmbert-small-nli, ONNX INT8

Dynamic-INT8 ONNX build of BalaRajesh1/mmbert-small-nli, whose upstream repo ships PyTorch weights only. Published so installs pull the quantized weights directly (141 MB instead of a 563 MB fp32 export) instead of re-exporting and re-quantizing per machine.

No retraining, no distillation: the same weights, exported and quantized.

Credits and license

• Fine-tuned NLI model: BalaRajesh1/mmbert-small-nli (MIT)
• Base encoder: jhu-clsp/mmBERT-small, Johns Hopkins (MIT)

MIT, inherited from both. All credit for the model belongs upstream; this repo only converts it.

Contents

The graph takes input_ids + attention_mask only, with no token_type_ids, so it runs on a plain ONNX Runtime session without a tokenizer type-id pass.

Why QUInt8 and not QInt8

Signed INT8 was not portable for this graph. Measured on a 72-pair labeled entailment set (AUC, entailed vs the rest), same fp32 export:

QInt8 was the best variant on one machine and worthless on another. Those two runs are not a controlled comparison (different CPU, and a different exporter/runtime produced each file), so the cause is not pinned down; a host without VNNI running signed INT8 through a saturating path is the prime suspect, and reduce_range=True did lift that host to 0.839. QUInt8 measured the same on both, so that is what ships.

If you quantize this model yourself, re-measure on your target host. An INT8 file that is excellent on one CPU can be a coin flip on another, and nothing in the file will tell you.

The tokenizer_class change

Upstream declares tokenizer_class: TokenizersBackend, which only Transformers v5 resolves: older Python Transformers raises, and transformers.js falls back to its base class with a warning. This repo declares PreTrainedTokenizerFast instead. Token ids are unchanged (verified identical on all 72 evaluation pairs); the algorithm lives in tokenizer.json, not in the class name.

Usage

import numpy as np, onnxruntime as ort
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("Nicolassuez/mmbert-small-nli-onnx-q8")
sess = ort.InferenceSession("onnx/model_quantized.onnx", providers=["CPUExecutionProvider"])

enc = tok(["The plant opened in 2019."], ["La station a ouvert en 2019."], return_tensors="np")
logits = sess.run(None, {k: v.astype(np.int64) for k, v in enc.items() if k != "token_type_ids"})[0]
p = np.exp(logits - logits.max(-1, keepdims=True))
print("P(entailment) =", (p / p.sum(-1, keepdims=True))[0][0])

Also loads in transformers.js with dtype: 'q8'.

Reported upstream accuracy

MNLI-matched 85.56 · SNLI test 88.27 · XNLI test (15 languages) 77.72 · WANLI test 69.18. These are the upstream fp32 figures, not re-measured here; the table above is the only measurement made on this build.