colbert-muvera-small-onnx

colbert-muvera-small-onnx

ONNX export of NeuML/colbert-muvera-small (Apache-2.0) for fastlate and any ONNX Runtime client. The weights are unchanged; this repo adds the graph, an INT8 dynamic quantization of it, and the ColBERT settings the model was trained with.

Files

• model.onnx: transformer, the PyLate Dense projection layer(s), and per-token L2 normalization, in one graph. Inputs input_ids, attention_mask (int64, dynamic batch and length); output embeddings of shape (batch, seq, 128). Opset 17.
• model_int8.onnx: onnxruntime.quantization.quantize_dynamic of the above, QInt8 weights.
• onnx_config.json: prefix ids, lengths, skiplist, expansion and padding settings copied from the PyLate model, plus fde_center.
• tokenizer.json: the source tokenizer, with [Q] /[D] as added tokens.

Settings

Query prefix [Q] (id 30522), document prefix [D] (id 30523), inserted right after the sequence-start token. Query length 32, document length 300. do_query_expansion=True: queries are padded to query_length with [MASK] (id 103), attention off on the padding, and every query vector is kept. Punctuation tokens (skiplist_words) are dropped from document embeddings after encoding, as in PyLate.

fde_center=False: whether to subtract the corpus mean token vector before building MUVERA fixed-dimensional encodings. Measured on a 3,000-chunk code/notebook corpus with FDEs of 4,096 dims: centring lowered shortlist recall for this model, so leave vectors as they are.

Validation

fp32 graph vs PyTorch forward: max abs diff 5.4e-07; int8 graph vs PyTorch: 0.122.

Exported with PyLate 1.6.0, transformers 5.3.0, torch 2.14.0, onnxruntime 1.2x, on 2026-09-18.