README.md

---
license: mit
language:
- code
- multilingual
tags:
- knowledge-graph
- relation-extraction
- relation-classification
- cross-encoder
- text-classification
- knowledge-distillation
- tensorrt
- openvino
pipeline_tag: text-classification
base_model:
- cross-encoder/mmarco-mMiniLMv2-L12-H384-v1
datasets:
- unicamp-dl/mmarco
---

# code-daemon-relation-v1

A **117M-parameter relation classifier**. Mark two entities inside a passage and it answers, in one
forward pass, **how they relate** — one of three relation types, or *no relation*.

It does a job usually handed to a large generative model — read a passage, extract typed relations
between the things it mentions — as **a single classification** instead of token-by-token generation.
That makes it cheap enough to sweep an entire corpus: **~2 900 pairs/sec** on a laptop RTX 5060.

```python
logits = session.run(None, {"input_ids": ids, "attention_mask": mask})[0]   # [B, 4]
```

Text is fed as an *(empty query, marked passage)* pair. Multilingual — the XLM-R backbone reads prose
and code comments in many languages.

---

## 1. The four classes

Wrap each entity in `[E1]…[/E1]` and `[E2]…[/E2]` inside its natural context. Take the `argmax`;
class 0 is an explicit **abstain**, and a softmax threshold drops the rest of the low-confidence tail.

| idx | label | meaning |
|--:|---|---|
| 0 | `NO_RELATION` | the two co-occur but are not related — **abstain** |
| 1 | `semantically_similar_to` | near-duplicate purpose, or the same goal by a different mechanism |
| 2 | `invalidates_with` | one supersedes, replaces or contradicts the other |
| 3 | `depends_on` | one requires or configures the other |

The taxonomy is deliberately **coarse**. An earlier 8-way version split these into near-synonym pairs
(`semantically_similar_to` vs `shares_purpose_with`, `replaced_by` vs `contradicts`, `depends_on` vs
`configured_by`) and the distinctions were not reliably separable from context — the classifier spent
its capacity on boundaries that downstream consumers then collapsed anyway. Merging them into three
positives plus abstain is what the model is actually good at.

**Decision rule as shipped:** `argmax != NO_RELATION` **and** `1 - softmax[NO_RELATION] >= tau`.
Gating on the *probability that any relation exists* rather than on the winning class's own
probability is more robust: when a real relation's mass spreads across two plausible classes, the
per-class maximum sags while "something is here" stays high.

---

## 2. Architecture

| | |
|---|---|
| Warm-start | [`cross-encoder/mmarco-mMiniLMv2-L12-H384-v1`](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1) |
| Encoder | XLM-RoBERTa, **12 layers / 384 hidden / 12 heads**, FFN 1536 |
| Vocabulary | **250 006** SentencePiece pieces = 250 002 XLM-R + 4 marker tokens |
| Markers | `[E1]` `[/E1]` `[E2]` `[/E2]` (ids 250002–250005) |
| Sequence | **256 tokens** on the shipped engines (64 / 128 / 320 also provided) |
| Inputs | `input_ids`, `attention_mask` — **no `token_type_ids`** |
| Output | `logits[batch, 4]` |
| Parameters | **~117M**, of which 96M is the multilingual embedding table |

### Entity-marker pooling, not `[CLS]`

The classification head does **not** read the `[CLS]` vector. It mean-pools the hidden states at the
**entity-start markers** — the `[E1]` and `[E2]` positions — concatenates the two, and passes that
through a single linear layer.

This matters for a relation task. A `[CLS]` vector summarises the whole passage, so the head has to
recover *which two things* the question is about from a global summary. Reading the marker positions
instead gives the head both arguments directly and in order, so the relation is scored between the
two entities rather than inferred from the sentence as a whole. Direction comes free: swap the
markers and the input genuinely changes.

---

## 3. How it was made

**Warm-started** from a strong multilingual ranking cross-encoder, with its single ranking logit
replaced by the 4-class marker-pooling head, then fine-tuned by **sequence-level distillation**.

A large instruction-tuned LLM (Claude) read real documentation and emitted relation tuples grounded
in the passage it was shown. Those tuples became the training targets after several filtering passes:

- **Grounding** — entity names the teacher marked must actually resolve inside the chunk, which drops
  hallucinated arguments.
- **Windowing** — the passage is cropped so that both markers survive truncation.
- **Merging** — the original 8 labels are collapsed to the 4 above.
- **Negatives** — explicit `NO_RELATION` examples synthesised from co-occurring but untupled entity
  pairs in the same chunk, so abstain is trained rather than inferred.
- **Logit adjustment** — the class prior is corrected at the loss, since a teacher naturally emits
  "similar" far more often than "depends on".

---

## 4. Speed

Measured on one laptop: Intel Core Ultra 9 275HX / NVIDIA RTX 5060 Laptop.

| lane | batch × seq | per batch | throughput | per pair |
|---|---|--:|--:|--:|
| **TensorRT FP16, RTX 5060 Laptop** | 16 × 256 | **5.44 ms** | **2 942 pairs/s** | 0.34 ms |
| **OpenVINO FP16, iGPU** (OV 2026.3) | 16 × 256 | **143 ms** | **111 pairs/s** | 8.96 ms |
| OpenVINO FP16, CPU (OV 2026.3) | 16 × 256 | 366 ms | 44 pairs/s | 22.9 ms |
| ONNX Runtime FP32, CPU | 16 × 256 | 380 ms | 42 pairs/s | 23.7 ms |

Per bucket, OpenVINO 2026.3 on the same laptop (pairs/s, solo):

| bucket | batch × seq | CPU FP16 | iGPU FP16 |
|---|---|--:|--:|
| s | 16 × 64 | 159 | **508** |
| m | 16 × 128 | 85 | **245** |
| l | 16 × 256 | 44 | **111** |
| xl | 16 × 320 | 34 | **85** |

The integrated GPU is ~2.6× the CPU on every bucket, so on a host without a discrete card the
iGPU lane is the one to route relation extraction to. Both devices at once give ~87 % of the sum
of their solo rates — they share one memory controller.

The GPU lane is ~67× the CPU lane, which is the point: relation extraction over a corpus means tens
of thousands of candidate pairs, and only the compiled-engine path makes that a background task
rather than a batch job.

Four **length buckets** ship — seq 64 / 128 / 256 / 320 at batch 16. Attention is quadratic in
sequence length, so routing short passages to a short engine is worth taking when your pairs vary in
length. Padding is attention-masked, so a pair produces the same logits from any bucket that fits it.

### At corpus scale

The table above is a per-batch micro-benchmark. Sweeping a real corpus batches far wider, which
changes what the bottleneck is:

| | measured |
|---|--:|
| GPU inference, 267 pairs per call at seq 256 | **5.87 ms** → ~45 000 pairs/s |
| End-to-end incl. tokenisation (12 CPU threads, overlapped) | **~2 500 pairs/s** |
| Relations written, 3 043 chunks | 15 175 in 15.3 s → **994/s** |
| Relations written, 4 975 chunks | 31 696 in 45.0 s → **704/s** |

Once the engine is batched this wide the GPU is no longer the limit — **tokenisation is**, and it is
worth giving it real thread count. Measured padding waste at these batch sizes is 18–20%, which is
what the length buckets are there to keep down.

---

## 5. Standalone use

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

LABELS = ["NO_RELATION", "semantically_similar_to", "invalidates_with", "depends_on"]

tok  = AutoTokenizer.from_pretrained(".")        # includes the [E1]/[E2] marker tokens
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

def classify(marked_text, max_len=256, tau=0.7):
    enc = tok([""], [marked_text], padding="max_length", 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][0]
    p = np.exp(logits - logits.max()); p /= p.sum()
    if 1.0 - p[0] < tau:                          # gate on "any relation at all"
        return "NO_RELATION", float(p[0])
    i = int(p.argmax())
    return (LABELS[i], float(p[i])) if i else ("NO_RELATION", float(p[0]))

classify("The [E1]FAISS[/E1] index was replaced by the [E2]native IVF[/E2] backend.")
# -> ('invalidates_with', 0.7x)
```

Both entities must appear **inside one passage**, marked in place. The model reads context, so a bare
pair of names with no surrounding text carries little signal.

---

## 6. Evaluation

**Dev macro-F1 = 0.547** over the four classes, on a held-out split of the distillation set.

Read that as what it is. The classes are intrinsically imbalanced — a teacher describing
documentation emits "similar" far more often than "depends on" — and the merged taxonomy still
contains genuinely ambiguous boundaries that human annotators would also disagree on. The abstain
class plus the `tau` gate exist because the useful operating point is *high-precision edges*, not
maximum recall: for building a graph, the real test is spot-checking the edges it emits at your
chosen threshold.

**Suited to**
- Turning prose or documentation into a typed concept graph.
- Any sweep where a large LLM per pair would be too slow or too expensive.
- Multilingual corpora, including code comments.

**Not suited to**
- Fine-grained relation ontologies — this is 3 positives plus abstain by design.
- Entity *extraction*: it classifies pairs you already found, it does not find them.
- Passages where the two entities are far apart — the marked window is 256 tokens.

---

## 7. What is in this repo

Compiled engines, named per **runtime × OS × GPU arch**, plus the ONNX for standalone use.

- **TensorRT FP16** — `code-daemon-relation-v1-{s,m,l,xl}_{win_x64,linux_x64}_trt11.0_sm_120.engine`
  (buckets seq 64 / 128 / 256 / 320, batch 16).
- **OpenVINO FP16** — `code-daemon-relation-v1-{s,m,l,xl}_ov2026.3_{cpu,igpu}_fp16_b16_s{64,128,256,320}.{xml,bin}`.
- **Tokenizer** — `tokenizer.json`, `sentencepiece.bpe.model`, `tokenizer_config.json` (XLM-R
  SentencePiece with the four marker tokens added).
- **ONNX** — `model.onnx` (+ `model.onnx.data`), FP32, the build source for every engine above.

FP16 rather than INT8: this architecture's activation outliers make per-tensor INT8 calibration
lossy, and FP16 costs nothing on any GPU that can run it.

---

## 8. License & attribution

Released **MIT**.

| Source | Note |
|---|---|
| `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` (warm-start) | **mMARCO ← MS MARCO → non-commercial research terms** |
| Distillation targets (LLM-labelled open-source documentation) | self-generated |
| Synthesised negatives and rare-class augmentation | generated |

⚠️ 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 matters to you.

Warm-started from **[cross-encoder/mmarco-mMiniLMv2-L12-H384-v1](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1)**.
Backbone: XLM-RoBERTa. Used by the [UltraCode](https://github.com/faxenoff/ultracode) code assistant,
though nothing about the model is specific to it.