| --- |
| license: mit |
| language: |
| - code |
| - multilingual |
| tags: |
| - code |
| - code-search |
| - code-retrieval |
| - embeddings |
| - feature-extraction |
| - sentence-similarity |
| - knowledge-distillation |
| pipeline_tag: feature-extraction |
| base_model: |
| - nomic-ai/CodeRankEmbed |
| 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** 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. |
|
|
| 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. |
|
|
| - **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**. |
|
|
| ## How it was made |
|
|
| 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). |
|
|
| 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. |
|
|
| ## Built for speed |
|
|
| This model trades long-context capability for raw throughput on short code units: |
|
|
| - **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]`). |
|
|
| ## Intended use |
|
|
| - **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. |
|
|
| 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: |
|
|
| ```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 |
|
|
| 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. |
|
|
| - **TensorRT** `*.engine` — NVIDIA, INT8 W8A16, per arch × OS × bucket: |
| `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. |
|
|
| ## Evaluation — in-scope CoIR (sub-CoIR) |
|
|
| 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): |
|
|
| | CoIR task (in-scope) | 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 per language (NL→code): python **91.96**, go 82.27, java 76.02, php 68.98, ruby 65.94, js 62.66. |
|
|
| > 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. |
|
|
| 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. |
|
|
| ## Performance (embeddings / sec) |
|
|
| | Backend | Hardware | Throughput | |
| |---|---|--:| |
| | TensorRT INT8 | NVIDIA RTX 5060 (sm_120) | **~20,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. |
| |
| ## 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: |
| |
| | 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 **[nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed)** (MIT). Backbone: XLM-RoBERTa (MIT). |
| |
