Model card: R2 (gte-modernbert distill) + reranker section

README.md

@@ -13,7 +13,7 @@ tags:
 - knowledge-distillation
 pipeline_tag: feature-extraction
 base_model:
-- nomic-ai/CodeRankEmbed
+- Alibaba-NLP/gte-modernbert-base
 datasets:
 - Fsoft-AIC/the-vault-function
 - unicamp-dl/mmarco
@@ -24,54 +24,59 @@ datasets:
 
 # code-daemon-embed-v1
 
-A small, fast **code embedding model** purpose-built to vectorize a **code graph** (functions, methods,
-doc-chunks) for on-device semantic code search. It ships with the
-[UltraCode](https://github.com/faxenoff/ultracode) MCP server, running as a TensorRT / TVM / OpenVINO / ONNX engine.
+A small, fast **code-embedding model** for semantic code search and code retrieval. It maps short code
+units — function and method bodies, signatures, docstrings — and short natural-language queries into a
+shared **768-dim** vector space.
 
-It is **deliberately specialized for short code units, not long documents** — long-text handling was
-intentionally dropped (max sequence **128 tokens**) to maximize embedding throughput. Code-graph nodes
-are short (entity names, signatures, doc-chunks); spending capacity and latency on a long-context path
-would only slow the hot path it never uses.
+It is **specialized for short code, not long documents**: the maximum sequence is **128 tokens**,
+trading long-context capability for high throughput and strong quality on short units.
 
-- **768-dim** embeddings, **Matryoshka (MRL)** truncatable to **512 / 256** with graceful decay.
-- **~54.5M params** — XLM-RoBERTa architecture, **4 layers / 768 hidden**, **code-only 32k SentencePiece vocab**.
-- **Mean pooling** baked into the graph — output is already pooled (`[batch, 768]`); just **L2-normalize**.
-- Trained at sequence length **128**; length buckets s/m/l = seq **40 / 64 / 128**.
+- **768-dim** embeddings, **Matryoshka (MRL)** — truncatable to **512 / 256** dims with graceful decay.
+- **~54.5M params** — XLM-RoBERTa architecture, **4 layers / 768 hidden**, a **code-oriented 32k SentencePiece vocab**.
+- **Mean pooling fused into the graph** — the output is already pooled (`[batch, 768]`); just **L2-normalize**.
+- Trained at sequence length **128** (length buckets s / m / l = seq **40 / 64 / 128**).
 
-## How it was made
+## Where it's good — and where it isn't
 
-Knowledge-distilled (embedding regression) from the teacher **[`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)**
-(MIT, 137M, a strong code retriever). The student is a fresh, shallow-wide XLM-R encoder trained
-from scratch on the teacher's passage embeddings over a ~32M-sample code + text corpus, with a
-custom 32k code-oriented SentencePiece vocabulary (syntax + identifier lexicon rather than prose).
+Measured on CoIR (NDCG@10, full corpora). Use it for:
 
-Why shallow-wide (4l/768h) + code vocab: on an internal code-search golden set this **beat** both a
-deeper 6-layer variant and the earlier 64k-prose-vocab cut — depth hurt, a code-tuned vocab and a
-wide body helped.
+**Strong:**
+- **Code → related code** — finding similar or duplicate implementations (its best relative area).
+- **Natural-language → code** — docstring or description → function (strongest on Python).
+- **Short question → code** — "how / where does X…".
+- **NL → SQL** and **NL instruction → code**.
 
-## Built for speed
+**Weak / out of scope:**
+- **Long documents** — hard 128-token cap; longer inputs are truncated. This is not a long-context retriever.
+- **Noisy / ambiguous NL→code** (hard, under-specified queries) — mid quality.
+- **General English prose** (medical / financial / news) — the code-specialized 32k vocab trades
+  general-text coverage for code. Multilingual text works as a fallback, not a specialty.
 
-This model trades long-context capability for raw throughput on short code units:
+Embed **queries and documents the same way** — no instruction prefix. For smaller indexes, truncate to
+**256** or **512** dims (MRL) before normalizing.
 
-- **Short context by design** — max **128 tokens**, no long-document path. Code-graph nodes are short
-  (entity names, signatures, doc-chunks), so the model and its engines are tuned only for that, avoiding
-  the cost of a wide dynamic shape range.
-- **Rectangular TensorRT profiles** — each length bucket is built with a *fixed* shape (min == opt == max),
-  not a dynamic range, so the autotuner locks one optimal kernel set per bucket:
-  **s** = batch 64 × seq 40 · **m** = batch 128 × seq 64 · **l** = batch 256 × seq 128.
-- **INT8 (W8A16)** weights; **mean-pool + projection + L2-norm fused into the graph** (one pass → `[B, 768]`).
+## How it was made
 
-## Intended use
+Knowledge-distilled (embedding regression) from **[`Alibaba-NLP/gte-modernbert-base`](https://huggingface.co/Alibaba-NLP/gte-modernbert-base)**
+(Apache-2.0, ~150M, a strong general + code retriever on a ModernBERT backbone). The student is a fresh,
+shallow-wide XLM-R encoder trained from scratch on the teacher's passage embeddings over a ~30M-sample
+code + text corpus, with a custom 32k code-oriented SentencePiece vocabulary (syntax + identifier lexicon
+rather than prose). The shallow-wide 4-layer / 768-hidden shape keeps inference cheap while distilling at
+the teacher's full 768-dim width.
+
+## Built for speed
+
+- **Short context by design** — max **128 tokens**, no long-document path, so the engines avoid a wide
+  dynamic shape range.
+- **Rectangular TensorRT profiles** — each length bucket is a *fixed* shape (min == opt == max), one
+  optimal kernel set per bucket: **s** = batch 64 × seq 40 · **m** = batch 128 × seq 64 ·
+  **l** = batch 256 × seq 128.
+- **INT8 (W8A16)** weights; **mean-pool + projection + L2-norm fused into the graph** (one pass → `[B, 768]`).
 
-- **Semantic code search / code retrieval**, and general (multilingual) text retrieval as a fallback.
-- Embed **queries and documents the same way** (no instruction prefix — the student was distilled on
-  passage embeddings, unlike the teacher whose prefix is query-only). Mean-pool → **L2-normalize**.
-- For smaller indexes, truncate to **256** or **512** dims (MRL) before normalizing.
+## Usage (standalone ONNX)
 
-The daemon runs the bundled engines directly (this repo is its CDN), but the FP32 `model.onnx` is
-**also bundled** for standalone use. The recipe below runs it with `onnxruntime`: tokenize with the
-bundled `sentencepiece.bpe.model`, run, and the pooled `[B,768]` is already produced — just
-L2-normalize:
+The FP32 `model.onnx` is bundled. Tokenize with the bundled `sentencepiece.bpe.model`, run, and the
+pooled `[B, 768]` output is already produced — just L2-normalize:
 
 ```python
 import onnxruntime as ort, sentencepiece as spm, numpy as np
@@ -91,71 +96,67 @@ def embed(texts, max_len=128, mrl_dim=768):
 
 ## What's in this repo — ready-to-run compiled engines
 
-This repo holds **pre-compiled, ready-to-run engines**, named per
-**runtime × GPU arch × OS × length-bucket** — grab the compiled model that matches your runtime and
-hardware and use it directly, with no compilation on your machine.
+Pre-compiled engines, named per **runtime × GPU arch × OS × length-bucket** — pick the one matching your
+runtime and hardware; no compilation needed.
 
-- **TensorRT** `*.engine` — NVIDIA, INT8 W8A16, per arch × OS × bucket:
+- **TensorRT** `*.engine` — NVIDIA, INT8 W8A16:
   `code-daemon-embed-v1-{s,m,l}_{win_x64,linux_x64}_trt_sm_{86,89,120}.engine`
   (sm_86 ≈ RTX 30xx / A-series · sm_89 ≈ RTX 40xx / L4 · sm_120 ≈ RTX 50xx).
 - **TVM** `*_tvm_vulkan.{dll,so}` — Vulkan fallback for non-TRT / older NVIDIA & other GPUs, per bucket.
 - **OpenVINO** `*.xml` + `*.bin` — Intel **CPU / iGPU / NPU**, per bucket.
 - **Metal** `*_tvm_metal.*` — Apple Silicon (macOS), per bucket.
-- **Tokenizer** — `sentencepiece.bpe.model` (the model's SentencePiece; specials baked at
-  pad=0 / unk=1 / bos=2 / eos=3, byte-fallback) + `tokenizer_config.json`. The daemon loads the SP directly.
-- **ONNX source** — `model.onnx` (+ `model.onnx.data`) FP32 and `model_int8qdt.onnx` (INT8 W8A16) — for
-  standalone `onnxruntime` / optimum use, and the source the engines are compiled from.
+- **Tokenizer** — `sentencepiece.bpe.model` (specials at pad=0 / unk=1 / bos=2 / eos=3, byte-fallback) +
+  `tokenizer_config.json`.
+- **ONNX source** — `model.onnx` (+ `model.onnx.data`) FP32 and `model_int8qdt.onnx` (INT8 W8A16).
 
-## Evaluation — in-scope CoIR (sub-CoIR)
+## Evaluation — CoIR (NDCG@10, full corpora)
 
-CoIR is a broad code-retrieval benchmark, but **4 of its 10 tasks are out of scope** for a code-graph
-search engine (code↔code translation, multi-turn dialogue, long problem-statements — the daemon never
-performs these). The honest, relevant view is the **in-scope subset** — the retrieval patterns this
-model is actually built for (NDCG@10, full corpora):
+Retrieval quality on the CoIR tasks that match this model's design (short code + retrieval). Four of
+CoIR's ten tasks — code↔code translation, multi-turn dialogue, long problem-statements — exceed the
+128-token / retrieval scope and are not shown.
 
-| CoIR task (in-scope) | NDCG@10 | Pattern |
+| CoIR task | NDCG@10 | Pattern |
 |---|--:|---|
-| codesearchnet (6-lang avg) | **74.64** | docstring / NL → code (the core path) |
-| stackoverflow-qa | 53.18 | short question → code |
-| synthetic-text2sql | 50.15 | NL → SQL |
-| codefeedback-st | 47.71 | NL instruction → code |
-| codesearchnet-ccr (6-lang avg) | 44.30 | code → related code (clone/dup) |
-| cosqa | 32.14 | NL question → code (noisy / hard) |
-| **In-scope average (sub-CoIR)** | **51.56** | |
+| codesearchnet (6-lang avg) | **73.17** | docstring / NL → code |
+| stackoverflow-qa | 62.70 | short question → code |
+| synthetic-text2sql | 61.32 | NL → SQL |
+| codesearchnet-ccr (6-lang avg) | **57.30** | code → related code |
+| codefeedback-st | 56.38 | NL instruction → code |
+| cosqa | 35.51 | NL question → code (noisy / hard) |
+| **Average** | **57.73** | |
 
-codesearchnet per language (NL→code): python **91.96**, go 82.27, java 76.02, php 68.98, ruby 65.94, js 62.66.
+Per language — codesearchnet (NL→code): python **88.82**, java 75.78, php 73.71, go 73.49, ruby 64.10, js 63.09.
+Per language — codesearchnet-ccr (code→code): ruby **65.67**, java 62.87, js 61.05, python 55.88, go 51.53, php 46.77.
 
-> The full 10-task official CoIR average (36.67) is dragged down by the 4 out-of-scope tasks and is not
-> representative of the real query mix. For scale, the 1.5B-class `bge-code-v1` scores 81.77 on full
-> CoIR — this is a **54.5M** model (27× smaller) tuned for one job.
+**Binary (1-bit sign) vectors** retain **94.4%** of the float NDCG before any rescore — the embeddings
+are clean by sign, so a Hamming (XOR + popcount) index over 1-bit codes gives a ~32× memory / search win
+at a small quality cost.
 
-On the daemon's own `search-gold` golden set (its real query distribution): **hit@5 0.692** — +80% over
-the retired v1.1 cut (0.385). Binary (1-bit) vectors retain ~91% of float NDCG before rescore.
+> For scale, the 1.5B-parameter `bge-code-v1` scores 81.77 on full CoIR — this is a **54.5M** model
+> (27× smaller) tuned for short-code retrieval.
 
 ## Performance (embeddings / sec)
 
 | Backend | Hardware | Throughput |
 |---|---|--:|
-| TensorRT INT8 | NVIDIA RTX 5060 (sm_120) | **~20,000 emb/s** |
+| TensorRT INT8 | NVIDIA RTX 5060 (sm_120) | **~24,000 emb/s** |
 | OpenVINO INT4 | Intel iGPU (Xe2, Lunar Lake) | ~580 emb/s |
 | OpenVINO INT4 | Intel NPU (NPU4) | ~574 emb/s |
 | OpenVINO INT8 | Intel CPU (Core Ultra) | ~375 emb/s |
 | OpenVINO — **all 3 in parallel** | iGPU + NPU + CPU concurrently | ~1,290 emb/s |
 
-The combined figure is **genuine concurrent multi-device execution**: three independent workers — one
-bound to each of the iGPU, NPU and CPU — embed different batches **at the same time**, and the
-throughputs add up. This is **not** OpenVINO's `AUTO` mode (which selects a *single* device per
-inference and never runs the three simultaneously); the daemon length-sorts inputs and fans the buckets
-across all three devices. TRT is infer throughput on the bucketed batch path; OV figures measured on a
-Core Ultra (Lunar Lake) laptop.
+The combined OpenVINO figure is genuine concurrent multi-device execution — three workers (iGPU, NPU,
+CPU) embed different batches at the same time and the throughputs add up. This is **not** OpenVINO's
+`AUTO` mode (which picks a single device per inference). Measured on a Core Ultra (Lunar Lake) laptop;
+the TensorRT figure is on the bucketed batch path.
 
 ## License & training data
 
 Released under the **MIT license**.
 
-The teacher (`nomic-ai/CodeRankEmbed`) is MIT, and the XLM-R architecture is MIT. As is standard
-practice for distilled embedding models, the **weights are released under MIT**. For transparency,
-the training corpus the teacher embedded includes:
+The teacher (`Alibaba-NLP/gte-modernbert-base`) is Apache-2.0, and the XLM-R architecture is MIT. As is
+standard practice for distilled embedding models, the **weights are released under MIT**. For
+transparency, the training corpus the teacher embedded includes:
 
 | Dataset | License note |
 |---|---|
@@ -166,10 +167,9 @@ the training corpus the teacher embedded includes:
 | `jinaai/negation-dataset` | see source repo |
 
 ⚠️ If your use requires strict training-data-license compliance, note that **mMARCO derives from
-MS MARCO (non-commercial)**. Whether a distilled model inherits dataset-use terms is legally
-unsettled; this is **not legal advice**. A data-clean variant can be retrained without the mMARCO
-splits if needed.
+MS MARCO (non-commercial)**. Whether a distilled model inherits dataset-use terms is legally unsettled;
+this is **not legal advice**. A data-clean variant can be retrained without the mMARCO splits if needed.
 
 ## Attribution
 
-Distilled from **[nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed)** (MIT). Backbone: XLM-RoBERTa (MIT).
+Distilled from **[Alibaba-NLP/gte-modernbert-base](https://huggingface.co/Alibaba-NLP/gte-modernbert-base)** (Apache-2.0). Backbone: XLM-RoBERTa (MIT).