--- license: mit language: - code - multilingual tags: - code - code-search - code-retrieval - reranker - cross-encoder - text-ranking - knowledge-distillation pipeline_tag: text-ranking base_model: - cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 datasets: - code_search_net - unicamp-dl/mmarco --- # 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](https://github.com/faxenoff/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`](https://huggingface.co/faxenoff/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`](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1) (a strong multilingual MS MARCO cross-encoder) and **fine-tuned with a listwise loss** — **ListNet top-1 softmax cross-entropy** over each query's *{1 positive + ≤8 hard negatives}* group. The hard negatives were mined by the code retriever [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed) from CoIR (NL→code, code-dominant 0.75 share): the documents the first stage confuses with the answer are exactly what the reranker must learn to push down. Trained on ~99k query groups, 2 epochs, on a single A100. **Why listwise (not pointwise BCE):** a listwise loss optimizes the *order of the whole candidate group*, not each pair in isolation, so it spends capacity on the **top of the list** — which is all a reranker is for. Against the same model trained with pointwise BCE, listwise lifted **Hit@1 +0.12 and MRR +0.10** on our golden set; recall (Hit@5/@10) is unchanged because that ceiling is set by the first-stage retriever, not the 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. ```python 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` + `tokenizer_config.json` (XLM-R SP; the daemon loads it directly). - **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 & training data Released under the **MIT license** (the mmarco base + XLM-R backbone are MIT/Apache; fine-tuned weights released MIT). Training-data transparency: | Source | Note | |---|---| | `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` (warm-start base) | **mMARCO ← MS MARCO → non-commercial research terms** | | CoIR (NL→code hard-negative mining) | code-retrieval corpora; mixed upstream provenance | | hard-neg miner `nomic-ai/CodeRankEmbed` | MIT | ⚠️ 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](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1)**. Hard negatives mined with **[nomic-ai/CodeRankEmbed](https://huggingface.co/nomic-ai/CodeRankEmbed)** (MIT). Backbone: XLM-RoBERTa.