code-daemon-embed-v1

A small, fast code-embedding model for semantic code search over real codebases. 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 real retrieval traffic, not benchmark prose: the training queries mirror how agents and developers actually search a repository (short keyword-bag and behavior queries), and the maximum sequence is 128 tokens, trading long-context capability for high throughput and strong quality on short units.

• 768-dim embeddings (full width; no MRL truncation in this release).
• ~75.7M params — XLM-RoBERTa architecture initialized from multilingual-e5-base with strided 12→8 layer truncation and the vocabulary pruned 250k → 22,739 pieces (English + code coverage kept; the transformer layers see identical tokenizations for EN/code text).
• A built-in reranker — see below: cross-encoder ranking knowledge is distilled into the embedding space, so the top of the ranking behaves like a reranked list at bi-encoder cost.
• 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).

The built-in reranker (why the top of the ranking is sharp)

Classic bi-encoders compress each document into a vector before seeing the query, so they cannot do the fine-grained pairwise comparison a cross-encoder reranker does. This model closes part of that gap at training time instead of inference time:

• Hard negatives were mined for every training query (documents a strong retriever ranks high but wrong).
• Qwen/Qwen3-Reranker-4B (a cross-encoder) scored every (query, candidate) pair with graded relevance logits.
• The student was trained with a listwise-KL objective to reproduce the reranker's ranking distribution — not just "positive above negatives", but how much each near-miss should trail.

The result: reranker-flavored top-rank precision with zero inference overhead — no second model, no second pass, no extra latency. In the production system this model was built for, enabling an actual runtime reranker on top of it measured as net-negative: the distillation had already captured the useful part.

Where it's good — and where it isn't

Strong (the design target):

• Real repository search — short NL/keyword queries by developers and coding agents over an indexed codebase ("git watcher head change reindex", "acquire database lock for project hash").
• Natural-language → code at short lengths; identifier and mixed queries.
• Complementing a lexical (BM25-style) channel in a hybrid retriever — it was trained for that role.

Weak / out of scope:

• Long documents — hard 128-token cap; not a long-context retriever.
• Benchmark-style long problem statements (competitive-programming prompts, multi-turn dialogue, code↔code translation) — deliberately out of scope, see the evaluation note.
• General English prose (medical / financial / news) — the pruned vocab and code specialization trade general-text quality away.

Embed queries and documents the same way — no instruction prefix.

How it was made

Three-teacher pipeline on a pretrained backbone (the full recipe and every ablation is documented in the source project):

1. Backbone: intfloat/multilingual-e5-base (12L XLM-R retriever), strided-truncated to 8 layers (keep {0,1,3,5,7,9,10,11}) with the embedding table pruned and re-indexed to a 22,739-piece SentencePiece vocabulary in a raw-SP convention (pad=0 / unk=1 / bos=2 / eos=3).
2. Dense representation distillation from Alibaba-NLP/gte-modernbert-base (MSE to the teacher's embeddings) over real repository entities in a front-loaded serve format (semantics before identifiers, within the 128-token budget).
3. Listwise-KL ranking distillation from Qwen/Qwen3-Reranker-4B cross-encoder scores over mined hard negatives (12 kept per query) — the "built-in reranker".
4. Diversity regularization: ~26% of the grouped stream is out-of-domain code-retrieval pairs (CoIR-derived). Measured to be the sweet spot: 0% loses deep recall, 50% dilutes the top. Counter-intuitively, the dissimilarity of these rows is the active ingredient.

Training queries were generated few-shot in the style of real agent search traffic (short keyword-bags, behavior descriptions, identifier fragments) with identifier-leak gating — not docstring paraphrases. This is the main reason benchmark numbers below understate real-repo performance.

Built for speed

• Short context by design — max 128 tokens; the engines avoid wide dynamic-shape ranges.
• Rectangular TensorRT profiles — each length bucket is a fixed shape (min == opt == max): s = batch 64 × seq 40 · m = batch 128 × seq 64 · l = batch 256 × seq 128.
• Mean-pool + L2-friendly head fused into the graph (one pass → [B, 768]).
• FP32/FP16 engines in this release (INT8 PTQ measurably hurts this checkpoint's top-rank precision; a QAT INT8 variant is planned).

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:

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):
    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 pooled [B,768]
    return out / np.linalg.norm(out, axis=1, keepdims=True)

What's in this repo

Named per runtime × GPU arch × OS × length-bucket (the filenames are a loading contract — do not rename):

• TensorRT *.engine — NVIDIA: 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 GPUs, per bucket.
• OpenVINO *.xml + *.bin — Intel CPU / iGPU / NPU, per bucket.
• Metal *_tvm_metal.* — Apple Silicon (macOS), per bucket.
• Tokenizer — sentencepiece.bpe.model (raw-SP: pad=0 / unk=1 / bos=2 / eos=3, byte-fallback) + tokenizer_config.json.
• ONNX source — model.onnx (+ model.onnx.data) FP32.

Compiled engines for this release are being rebuilt and uploaded progressively; the ONNX + tokenizer are the always-present source of truth.

Evaluation

Real-repository retrieval (the design target)

On a benchmark of real agent search queries over a live indexed codebase (80 captured queries, multi-positive, file-level), the production hybrid retriever built on this model reaches:

metric	value
hit@5	0.73
hit@1	0.44
pool_recall@64	0.79

This is the metric the model is optimized for — short real-world queries against a real repository.

CoIR (NDCG@10, full corpora) — out-of-domain reference

CoIR queries are largely docstrings and long problem statements — a different style from real repo search traffic — so treat these as an out-of-domain lower bound: on real queries in real repositories the model performs substantially better than its CoIR positioning suggests (see above). 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. (~2k CoIR-derived rows were present in training as diversity regularization; overlap with test splits was not audited.)

CoIR task	NDCG@10	Pattern
synthetic-text2sql	62.59	NL → SQL
stackoverflow-qa	50.82	short question → code
codefeedback-st	47.42	NL instruction → code
codesearchnet (6-lang avg)	42.68	docstring / NL → code
codesearchnet-ccr (6-lang avg)	40.95	code → related code
cosqa	20.75	NL question → code (noisy / hard)
Average	44.20

Per language — codesearchnet (NL→code): python 66.80, go 51.30, java 38.77, php 38.16, js 32.17, ruby 28.90. Per language — codesearchnet-ccr (code→code): js 47.72, ruby 45.28, java 42.05, python 38.43, php 36.60, go 35.64.

Performance (embeddings / sec)

Being re-measured for this release's engines — table will be updated with the engine upload.

License & training data

Released under the MIT license.

The backbone (intfloat/multilingual-e5-base) is MIT; the teachers (gte-modernbert-base, Qwen3-Reranker-4B) are Apache-2.0. As is standard practice for distilled embedding models, the weights are released under MIT. Training corpus, for transparency:

Data	License note
Repository entity texts + synthetic search queries over an MIT-licensed codebase	MIT
CoIR-derived code-retrieval pairs (~26% diversity stream)	per-subset CoIR licenses (research benchmark)
Teacher embeddings / scores (gte-modernbert, Qwen3-Reranker-4B)	Apache-2.0 teachers

Unlike the previous release, no MS MARCO / mMARCO data was used — the earlier non-commercial caveat no longer applies. This is not legal advice.

Attribution

Backbone: intfloat/multilingual-e5-base (MIT). Dense teacher: Alibaba-NLP/gte-modernbert-base (Apache-2.0). Ranking teacher: Qwen/Qwen3-Reranker-4B (Apache-2.0).