| --- |
| license: mit |
| language: |
| - code |
| - multilingual |
| tags: |
| - code |
| - code-search |
| - code-retrieval |
| - embeddings |
| - feature-extraction |
| - sentence-similarity |
| - knowledge-distillation |
| pipeline_tag: feature-extraction |
| base_model: |
| - Alibaba-NLP/gte-modernbert-base |
| datasets: |
| - Fsoft-AIC/the-vault-function |
| - unicamp-dl/mmarco |
| - sentence-transformers/all-nli |
| - sentence-transformers/gooaq |
| - jinaai/negation-dataset |
| --- |
| |
| # code-daemon-embed-v1 |
|
|
| 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 **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** 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**). |
|
|
| ## Where it's good — and where it isn't |
|
|
| Measured on CoIR (NDCG@10, full corpora). Use it for: |
|
|
| **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**. |
|
|
| **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. |
|
|
| Embed **queries and documents the same way** — no instruction prefix. For smaller indexes, truncate to |
| **256** or **512** dims (MRL) before normalizing. |
|
|
| ## How it was made |
|
|
| 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]`). |
|
|
| ## Usage (standalone ONNX) |
|
|
| 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 |
| |
| sp = spm.SentencePieceProcessor(model_file="sentencepiece.bpe.model") # pad=0 unk=1 bos=2 eos=3 |
| sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"]) |
| |
| def embed(texts, max_len=128, mrl_dim=768): |
| ids = [[2, *sp.encode(t)[: max_len - 2], 3] for t in texts] # bos … eos |
| L = max(len(x) for x in ids) |
| inp = np.array([x + [0] * (L - len(x)) for x in ids], dtype=np.int64) # pad=0 |
| mask = (inp != 0).astype(np.int64) |
| out = sess.run(None, {"input_ids": inp, "attention_mask": mask})[0] # already mean-pooled [B,768] |
| out = out[:, :mrl_dim] # MRL truncation (768/512/256) |
| return out / np.linalg.norm(out, axis=1, keepdims=True) |
| ``` |
|
|
| ## What's in this repo — ready-to-run compiled engines |
|
|
| 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: |
| `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` (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 — CoIR (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 | NDCG@10 | Pattern | |
| |---|--:|---| |
| | 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** | | |
|
|
| 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. |
|
|
| **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. |
|
|
| > 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) | **~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 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 (`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 | |
| |---|---| |
| | `Fsoft-AIC/the-vault-function` (code) | dataset MIT; underlying code has mixed upstream provenance | |
| | `unicamp-dl/mmarco` (EN/RU retrieval) | **MS MARCO-derived → non-commercial research terms** | |
| | `sentence-transformers/all-nli` | SNLI (CC BY-SA 4.0) + MultiNLI | |
| | `sentence-transformers/gooaq` | Apache-2.0 | |
| | `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. |
| |
| ## Attribution |
| |
| Distilled from **[Alibaba-NLP/gte-modernbert-base](https://huggingface.co/Alibaba-NLP/gte-modernbert-base)** (Apache-2.0). Backbone: XLM-RoBERTa (MIT). |
| |
