mrsladoje
/

CodeRankEmbed-onnx-int8

@@ -1,124 +1,115 @@
----
-license: mit
-base_model: nomic-ai/CodeRankEmbed
-base_model_relation: quantized
-tags:
-- code
-- embeddings
-- onnx
-- int8
-- quantized
-language:
-- code
-pipeline_tag: feature-extraction
----
-# CodeRankEmbed — Dynamic INT8 Quantized (ONNX)
-A dynamically quantized INT8 version of [nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed), converted to ONNX by [jalipalo](https://huggingface.co/jalipalo/CodeRankEmbed-onnx) and quantized for fast CPU inference.
-## What is this?
-CodeRankEmbed is a 137M-parameter embedding model trained specifically for code search and retrieval. This repository provides a **dynamic INT8 weight-quantized** version that is significantly smaller and faster with negligible quality loss:
-| | FP32 (original) | INT8 (this model) |
-|---|---|---|
-| **File size** | 522 MB | 132 MB (−75%) |
-| **CPU inference** | 1.00× | ~2.09× faster |
-| **Min cosine vs FP32** | 1.000 | 0.961 |
-| **Calibration data needed** | — | None |
-Quantization was done with ONNX Runtime's `quantize_dynamic` (weights only, `QInt8`, `per_channel=True`). Activations remain in FP32 at runtime — the recommended approach for transformer/embedding models per the [ONNX Runtime documentation](https://onnxruntime.ai/docs/performance/model-optimizations/quantization.html).
-## Usage
-### With `@huggingface/transformers` (JavaScript / Node.js)
-```js
-import { pipeline } from "@huggingface/transformers";
-const extractor = await pipeline(
-  "feature-extraction",
-  "mrsladoje/CodeRankEmbed-onnx-int8",
-  { quantized: true }   // loads onnx/model_quantized.onnx automatically
-);
-const output = await extractor("def hello(): return 42", {
-  pooling: "mean",
-  normalize: true,
-});
-console.log(output.data); // Float32Array of 768 dimensions
-```
-### With `optimum` (Python)
-```python
-from optimum.onnxruntime import ORTModelForFeatureExtraction
-from transformers import AutoTokenizer
-model = ORTModelForFeatureExtraction.from_pretrained(
-    "mrsladoje/CodeRankEmbed-onnx-int8",
-    file_name="onnx/model_quantized.onnx",
-)
-tokenizer = AutoTokenizer.from_pretrained("mrsladoje/CodeRankEmbed-onnx-int8")
-inputs = tokenizer("def hello(): return 42", return_tensors="pt")
-outputs = model(**inputs)
-embeddings = outputs.last_hidden_state.mean(dim=1)
-```
-### With `onnxruntime` directly (Python)
-```python
-import onnxruntime as ort
-import numpy as np
-from tokenizers import Tokenizer
-tokenizer = Tokenizer.from_pretrained("mrsladoje/CodeRankEmbed-onnx-int8")
-tokenizer.enable_padding(length=128, pad_id=0)
-tokenizer.enable_truncation(max_length=128)
-session = ort.InferenceSession("onnx/model_quantized.onnx")
-encoded = tokenizer.encode("def hello(): return 42")
-input_ids = np.array([encoded.ids], dtype=np.int64)
-attention_mask = np.array([encoded.attention_mask], dtype=np.int64)
-outputs = session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
-embedding = outputs[1]  # sentence_embedding output, shape (1, 768)
-```
-## Quantization Details
-| Parameter | Value |
-|---|---|
-| Method | `quantize_dynamic` (ONNX Runtime) |
-| Weight type | `QInt8` (signed 8-bit integer) |
-| Scope | Weights only — activations quantized dynamically at runtime |
-| Per-channel | Yes |
-| Calibration | None required |
-| ORT version | 1.21.x |
-**Why dynamic over static?** Static INT8 quantization requires calibration data to pre-compute activation ranges. For transformer embedding models, activation distributions vary widely with input content and sequence length, making static calibration brittle (we validated this — static QDQ produced cosine similarities as low as 0.09–0.26 with MinMax calibration). Dynamic quantization sidesteps this entirely: weights are quantized offline and activations are quantized at runtime, giving robust quality across all inputs.
-## Quality Validation
-Validated on 10 code snippets across Python, JavaScript, Go, Java, Rust, TypeScript, and SQL:
-```
-Model            Size      Speedup    Min cosine vs FP32    Quality
-FP32 (baseline)  522.3 MB  1.00×      —                     baseline
-Dynamic INT8     132.2 MB  2.09×      0.9610                 excellent
-```
-A cosine similarity ≥ 0.96 means the INT8 embeddings point in essentially the same direction as FP32. For retrieval tasks — especially with a reranker in the pipeline — this difference is undetectable in practice.
-The ~2× CPU speedup is real compute acceleration (not just faster file loading), coming from ONNX Runtime's `MatMulIntegerToFloat` fused kernels operating on INT8 weights. VNNI-capable CPUs (Intel 10th gen+, AMD Zen4+) may see even larger gains.
-## Attribution
-- **Original model:** [nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed) — MIT License
-- **ONNX conversion:** [jalipalo/CodeRankEmbed-onnx](https://huggingface.co/jalipalo/CodeRankEmbed-onnx) — MIT License (inherited)
-- **INT8 quantization:** this repository — MIT License
-All work in this repository respects and complies with the MIT license of the original model.

+---
+library_name: transformers
+tags:
+- onnx
+- quantized
+- int8
+- code-search
+- embedding
+- nomic-bert
+base_model: nomic-ai/CodeRankEmbed
+license: mit
+---
+# CodeRankEmbed-onnx-int8
+INT8 quantized ONNX version of [nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed)
+for code search and embedding.
+## Quantization
+Dynamic INT8 quantization with **`reduce_range=True`** for cross-platform
+correctness:
+```python
+from onnxruntime.quantization import quantize_dynamic, QuantType
+quantize_dynamic(
+    model_input='CodeRankEmbed_fp32.onnx',
+    model_output='CodeRankEmbed_int8.onnx',
+    weight_type=QuantType.QInt8,
+    per_channel=True,
+    reduce_range=True,   # clamp weights to [-64, 63] for AVX2 kernel safety
+)
+```
+### Why `reduce_range=True`
+ORT's CPU INT8 MatMul kernels have two paths on x86:
+| CPU | Path | Full-range INT8 weights |
+|---|---|---|
+| Intel Cascade Lake+ / Ice Lake+ (VNNI) | `VPDPBUSD` | ✓ correct |
+| AMD Zen 4+ (VNNI / Genoa+) | `VPDPBUSD` | ✓ correct |
+| Apple Silicon (arm64 NEON + AMX) | separate arm64 kernels | ✓ correct |
+| **Intel pre-2019 / AMD Zen 3 Milan (AVX2 only)** | `pmaddubsw + phaddsw + paddd` | **✗ int16 accumulator overflows → degenerate output** |
+`reduce_range=True` clamps weights to `[-64, 63]` (7-bit signed range), giving
+the AVX2 `int16` intermediate enough headroom to avoid overflow. VNNI and arm64
+paths are unaffected (they handle full-range INT8 natively).
+### Known issue with earlier quantization
+A previous version of this model was quantized **without** `reduce_range=True`.
+It worked correctly on VNNI-capable CPUs and Apple Silicon, but produced
+degenerate embeddings (all texts mapping to near-identical vectors) on
+**AMD Zen 3 EPYC** and similar pre-VNNI x86 hosts — verified on RunPod
+RTX 5090 pods with EPYC 7543. This version fixes that. See commit history.
+## Performance
+- **Size**: 139 MB (FP32 source: 548 MB) — **~75% reduction**
+- **Output dim**: 768
+- **Expected cosine vs FP32**: ≥ 0.96 on production inputs
+- **Inference speedup (VNNI CPUs)**: ~2× vs FP32
+- **Inference speedup (pre-VNNI CPUs)**: ~1.5× vs FP32 (smaller win, but correct)
+### Validation (Mac M3 Max, ORT 1.24.3, 4-text probe)
+```
+T0 "how to parse json in python"        T3 "parse json data python"       cos=0.7749  (similar)
+T0 "how to parse json in python"        T2 "sql inner join three tables"  cos=0.1123  (dissimilar)
+Semantic separation                                                        0.6626  (≥ 0.15 healthy)
+```
+## Usage
+```python
+import onnxruntime as ort
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("mrsladoje/CodeRankEmbed-onnx-int8")
+session = ort.InferenceSession("model.onnx")
+inputs = tokenizer(
+    "your code or query here",
+    padding=True, truncation=True, max_length=512, return_tensors="np"
+)
+outputs = session.run(None, dict(inputs))
+# sentence_embedding is typically the second output; it's 768-dim L2-normalized
+```
+## Files
+- `onnx/model.onnx` — INT8 quantized model (139 MB)
+- `tokenizer.json`, `vocab.txt`, `config.json`, `special_tokens_map.json`, `tokenizer_config.json`
+  — from the base nomic-ai/CodeRankEmbed distribution
+## SHA256 (v2 — with `reduce_range=True`)
+```
+4eae31d09b1843103a1ebd5e2b2e24b5a5cad441a33906b35b12b1e2ed91d1db
+```
+Pin this in your downloader to guarantee you got the corrected weights and not
+a stale cached copy of v1.
+## Base model
+[nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed) (137M params),
+based on [Snowflake/snowflake-arctic-embed-m-long](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-long).
+ONNX conversion derived from [jalipalo/CodeRankEmbed-onnx](https://huggingface.co/jalipalo/CodeRankEmbed-onnx).
+## License
+MIT (inherited from base model).