--- license: apache-2.0 language: - en base_model: - zeroentropy/zerank-1-small pipeline_tag: text-ranking tags: - reranking - onnx - quantized - fastembed library_name: fastembed --- # zerank-1-small — ONNX Export ONNX export of [zeroentropy/zerank-1-small](https://huggingface.co/zeroentropy/zerank-1-small), a 1.7B Qwen3-based reranker. Includes three quantization levels for CPU inference. ## Files | File | Format | Size | Description | |------|--------|------|-------------| | `model.onnx` + `model.onnx_data` | FP16 | ~3.2 GB | Full precision | | `model_int8.onnx` + `model_int8.onnx_data` | INT8 | ~2.5 GB | Weight-only INT8 (per-tensor symmetric) | | `model_int4_full.onnx` | INT4 | ~1.3 GB | MatMulNBits INT4, block_size=32 | Conversion scripts: `export_zerank_v2.py` (FP16 export with dynamic batch), `stream_int8.py` (INT8 quantization). ## ⚠️ 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`:** ```python # 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) ```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 ```rust 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): | Task | Embedding only | cohere-rerank-v3.5 | Llama-rank-v1 | **zerank-1-small** | zerank-1 | |------|---------------|-------------------|--------------|----------------|----------| | Code | 0.678 | 0.724 | 0.694 | **0.730** | 0.754 | | Finance | 0.839 | 0.824 | 0.828 | **0.861** | 0.894 | | Legal | 0.703 | 0.804 | 0.767 | **0.817** | 0.821 | | Medical | 0.619 | 0.750 | 0.719 | **0.773** | 0.796 | | STEM | 0.401 | 0.510 | 0.595 | **0.680** | 0.694 | | Conversational | 0.250 | 0.571 | 0.484 | **0.556** | 0.596 | See [zeroentropy/zerank-1-small](https://huggingface.co/zeroentropy/zerank-1-small) for full details and Apache-2.0 license.