code-daemon-reranker-v1

A small, fast cross-encoder reranker purpose-built to re-order first-stage code-search hits for precision. It ships with the UltraCode MCP server as a TensorRT / OpenVINO / TVM engine, scoring (query, candidate) pairs after the embedding retriever has fetched a candidate pool.

A reranker is the second stage: the bi-encoder embed model (code-daemon-embed-v1) retrieves a pool fast, then this cross-encoder reads each (query, code) pair jointly and emits a single relevance logit, pulling the best match to the top. Joint attention over the pair is far more precise than the cosine of two independent vectors — at the cost of one forward pass per candidate, so it scores only a bounded pool (~64), not the whole index.

• ~117M params — XLM-RoBERTa 12 layers / 384 hidden, 250k multilingual SentencePiece vocab (the embedding table dominates the size).
• 2-input ONNX (input_ids, attention_mask; no token_type_ids) → a single relevance logit.
• Max sequence 256 tokens for the concatenated (query, document) pair.
• Listwise-trained — the key quality lever (below).

How it was made

Warm-started from cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 (a strong multilingual ranking cross-encoder) and fine-tuned with a listwise objective against groups of one positive and several hard negatives.

Why that matters for how you use it: a listwise objective optimizes the order of the whole candidate group rather than each pair in isolation, so the model spends its capacity on the top of the list. Expect the gains at Hit@1 / MRR / nDCG, and expect recall at depth to be unchanged — that ceiling belongs to the first-stage retriever, not to a reranker.

Built for speed

• Short context (256) — a (query, code-unit) pair is short; there is no long-document path.
• Bounded pool — the daemon reranks only the top ~64 fused candidates (the relevant file often sits at fused rank 30–60), then cuts back to top-k after reranking.
• Runs on the iGPU. On a box with both NVIDIA + Intel, the daemon routes the reranker to the Intel iGPU (OpenVINO fp16) and keeps CUDA free for the embedding model — the two search stages run on different devices in parallel. TensorRT (NVIDIA) and TVM/Vulkan engines are bundled for boxes without an Intel iGPU.
• fp16 — mmarco-format INT8 on the iGPU hits a known OpenVINO AccessViolation, so fp16 is shipped.

Intended use

Re-rank a candidate pool from a first-stage retriever for NL→code search (multilingual text works too). Feed (query, candidate) pairs, take the logit, sort descending. The score is a raw, unbounded logit (Identity head) — compare relatively within a query, not against a fixed threshold.

import onnxruntime as ort, numpy as np
from transformers import AutoTokenizer

tok  = AutoTokenizer.from_pretrained(".")            # bundled XLM-R SentencePiece
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

def rerank(query, docs, max_len=256):
    enc = tok([query] * len(docs), docs, padding=True, truncation=True,
              max_length=max_len, return_tensors="np", return_token_type_ids=False)
    logits = sess.run(None, {"input_ids":      enc["input_ids"].astype(np.int64),
                             "attention_mask": enc["attention_mask"].astype(np.int64)})[0]
    return sorted(zip(logits.reshape(-1).tolist(), docs), reverse=True)   # higher = more relevant

What's in this repo — ready-to-run compiled engines

Named per runtime × GPU arch × OS (single-profile — no length buckets):

• TensorRT code-daemon-reranker-v1_{win_x64,linux_x64}_trt_sm_{86,89,120}.engine — NVIDIA, fp16 (sm_86 ≈ RTX 30xx / A-series · sm_89 ≈ RTX 40xx / L4 · sm_120 ≈ RTX 50xx).
• OpenVINO code-daemon-reranker-v1_ov_{cpu,igpu}_fp16_b16_s256.{xml,bin} — Intel CPU / iGPU.
• TVM code-daemon-reranker-v1_*_tvm_vulkan.{dll,so} — Vulkan fallback for non-TRT / other GPUs.
• Tokenizer — sentencepiece.bpe.model (XLM-R SentencePiece; the daemon loads it directly, and AutoTokenizer.from_pretrained reads it too).
• ONNX source — model.onnx FP32 (the build source + standalone onnxruntime / optimum use).

Evaluation

Measured on the daemon's own search-gold golden set (26 NL→code queries — its real query distribution), reranking the embed retriever's top-64 pool. Metrics are advisory; manual review of the failures is the source of truth.

metric	embed-only	+ reranker	Δ
Hit@1	0.42	0.54	+0.12
Hit@3	0.62	0.77	+0.15
Hit@5	0.73	0.77	+0.04
Hit@10	0.81	0.81	0 (recall ceiling — set by the retriever)
MRR@10	0.55	0.65	+0.10
nDCG@10	0.59	0.67	+0.08

The gains concentrate on the top of the list (Hit@1, MRR, nDCG) — the listwise signature. Hit@10 is flat because a reranker can only reorder what the retriever already fetched.

Performance

Backend	Hardware	Rerank latency
OpenVINO fp16	Intel iGPU (Xe)	~550 ms / 64-candidate pool (~8.6 ms/candidate)

The daemon runs the reranker on the iGPU so the NVIDIA GPU stays free for the embedding model; the TensorRT path is faster per pass where an NVIDIA GPU is used for reranking.

License

Released under the MIT license (the warm-start base + XLM-R backbone are MIT/Apache; the fine-tuned weights are released MIT).

Source	Note
cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 (warm-start base)	mMARCO ← MS MARCO → non-commercial research terms

⚠️ The warm-start base derives from MS MARCO (non-commercial). Whether a fine-tuned model inherits dataset-use terms is legally unsettled; this is not legal advice. Retrain from a permissive base if strict compliance is required.

Attribution

Warm-started from cross-encoder/mmarco-mMiniLMv2-L12-H384-v1. Backbone: XLM-RoBERTa.