README.md

---
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.