zerank-1-small-ONNX

zerank-1-small — ONNX Export

ONNX export of zeroentropy/zerank-1-small, a 1.7B Qwen3-based reranker. Includes three quantization levels for CPU inference.

Files

Conversion scripts: export_zerank_v2.py (FP32 export with dynamic batch), stream_int8.py (INT8 quantization).

Note on the batch-broadcast patch (2026-05-03)

The original FP32 export and downstream INT8/INT4 derivatives all had a hardcoded batch=1 in the attention-mask And kernel: ONNX Runtime would crash at any batch>1 with Shape mismatch attempting to re-use buffer. {1,1,T,T} != {B,1,T,T}. The patch inserts Expand + Shape nodes before the And so the mask broadcasts dynamically. INT4 was re-uploaded on 2026-05-03 with the patch applied; the original unpatched INT4 is preserved at model_int4_full.unpatched.onnx for archival. Validated under fastembed-rs' reranker_parity harness — Spearman 0.99 vs FP32 reference, top-1 match across the 4 reference query groups.

⚠️ Important: chat template required

This model is a Qwen3-based causal LM that scores (query, document) relevance by extracting the "Yes" token logit at the last position. It requires a specific prompt format — plain pair tokenization produces meaningless scores.

Always format inputs using the Qwen3 chat template with system=query, user=document:

# using the tokenizer directly (matches training format exactly):
messages = [
    {"role": "system", "content": query},
    {"role": "user",   "content": document},
]
text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

This produces the following fixed string (equivalent, usable without a tokenizer):

<|im_start|>system
{query}
<|im_end|>
<|im_start|>user
{document}
<|im_end|>
<|im_start|>assistant

Usage with ONNX Runtime (Python)

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

MODEL_PATH = "model_int8.onnx"   # or model.onnx, model_int4_full.onnx
MAX_LENGTH = 512

sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
tok  = AutoTokenizer.from_pretrained("cstr/zerank-1-small-ONNX")

def format_pair(query: str, doc: str) -> str:
    messages = [
        {"role": "system", "content": query},
        {"role": "user",   "content": doc},
    ]
    return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

def rerank(query: str, documents: list[str]) -> list[float]:
    scores = []
    for doc in documents:
        text = format_pair(query, doc)
        enc  = tok(text, return_tensors="np", truncation=True, max_length=MAX_LENGTH)
        logit = sess.run(["logits"], {
            "input_ids":      enc["input_ids"].astype(np.int64),
            "attention_mask": enc["attention_mask"].astype(np.int64),
        })[0]
        scores.append(float(logit[0, 0]))
    return scores

query = "What is a panda?"
docs  = [
    "The giant panda is a bear species endemic to China.",
    "The sky is blue and the grass is green.",
    "Pandas are mammals in the family Ursidae.",
]
scores = rerank(query, docs)
for s, d in sorted(zip(scores, docs), reverse=True):
    print(f"[{s:.3f}] {d}")
# [+6.8] The giant panda is a bear species endemic to China.
# [+2.1] Pandas are mammals in the family Ursidae.
# [-5.8] The sky is blue and the grass is green.

Batch inference: The v2 export (model.onnx) supports batch_size > 1 via a dynamic causal+padding mask. Pad a batch with the tokenizer and pass the full batch at once for higher throughput.

Usage with fastembed-rs

use fastembed::{RerankInitOptions, RerankerModel, TextRerank};

let mut reranker = TextRerank::try_new(
    RerankInitOptions::new(RerankerModel::ZerankSmallInt8)
).unwrap();

// The chat template is applied automatically; batch_size > 1 is supported.
let results = reranker.rerank(
    "What is a panda?",
    vec![
        "The giant panda is a bear species endemic to China.",
        "The sky is blue.",
        "Pandas are mammals in the family Ursidae.",
    ],
    true,
    Some(32),
).unwrap();

for r in &results {
    println!("[{:.3}] {}", r.score, r.document.as_ref().unwrap());
}

Export details

export_zerank_v2.py wraps Qwen3ForCausalLM in a ZeRankScorerV2 that:

1. Builds a 4D causal+padding attention mask explicitly from input_ids.shape[0] — this makes the batch dimension dynamic in the ONNX graph (enabling batch_size > 1).
2. Runs the transformer body → hidden [batch, seq, hidden]
3. Gathers the hidden state at the last real-token position (attention_mask.sum - 1)
4. Applies lm_head, slices the "Yes" token (id 9454) → [batch, 1]

Output: logits [batch, 1] — raw Yes-token logit (higher = more relevant). FP16 weights, opset 18.

stream_int8.py performs fully streaming weight-only INT8 quantization:

• Never loads the full 6.4 GB FP32 model into RAM (peak ~1.5 GB)
• Symmetric per-tensor quantization: scale = max(|w|) / 127
• Adds DequantizeLinear → MatMul nodes for all MatMul B-weights
• Non-MatMul tensors (embeddings, LayerNorm) kept as FP32

Benchmarks (from original model card)

NDCG@10 with text-embedding-3-small as initial retriever (Top 100 candidates):

See zeroentropy/zerank-1-small for full details and Apache-2.0 license.

Provenance and EU AI Act Art. 53 note

• Upstream model: zeroentropy/zerank-1-small — published by zeroentropy.
• Upstream licence: apache-2.0. This repository redistributes under the same terms; it grants no rights the upstream licence does not.
• What was done here: format conversion and/or quantisation only (ONNX). No training, no fine-tuning, no merging, no distillation, no change to architecture, vocabulary or capability. Only the numeric representation of the upstream weights differs.
• Training data: documented — where it is documented at all — by the upstream provider; see the upstream model card. No training data was used, added or selected by this repository. No training-content summary was found on the upstream model card at the time of writing; that documentation gap is upstream's and is not filled here.
• Provider status: under Regulation (EU) 2024/1689 the upstream authors remain the provider of this model. Converting the serialisation format does not make this repository the provider of a new general-purpose AI model, and no such claim is made. Questions about training content, copyright policy or model capability belong upstream.