--- license: mit language: - code - multilingual tags: - code - knowledge-graph - relation-extraction - cross-encoder - text-classification - knowledge-distillation pipeline_tag: text-classification base_model: - cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 datasets: - unicamp-dl/mmarco --- # code-daemon-relation-v1 A tiny **relation classifier**: given two entities marked inside a piece of text, it predicts **how they relate** — one of 8 relation types, or *no relation*. It reads the two entities and their surrounding context **jointly** (a cross-encoder) and emits 8 class logits in a single forward pass. The point of the model is to do a job people usually hand to a large generative LLM — reading a passage and extracting typed relations between the things it mentions — as **one cheap classification pass** instead of token-by-token generation. It was **distilled from a 7B instruct model** into a ~117M cross-encoder, so it runs on a CPU/iGPU and is fast enough to sweep thousands of entity pairs when building a knowledge graph. It is used by the [UltraCode](https://github.com/faxenoff/ultracode) code assistant to turn documentation into a graph of related concepts, but nothing about it is specific to that tool. - **~117M params** — XLM-RoBERTa **12 layers / 384 hidden**, 250k **multilingual** SentencePiece vocab + 4 entity-marker tokens (`[E1] [/E1] [E2] [/E2]`), so the embedding table is 250 006 rows. - **2-input ONNX** (`input_ids`, `attention_mask`; no `token_type_ids`) → **`logits[batch, 8]`**. - **Max sequence 256** tokens for the marked passage. - **Multilingual** — the XLM-R backbone handles text and code comments in many languages. ## The 8 relation classes You mark the two entities with `[E1]…[/E1]` and `[E2]…[/E2]` in their context; the model returns a logit per class. Take `argmax`; class 0 (`NO_RELATION`) is the **abstain** class, and a softmax-probability threshold lets you drop low-confidence pairs. | idx | label | meaning | |--:|---|---| | 0 | `NO_RELATION` | the two entities co-occur but are not related (abstain) | | 1 | `semantically_similar_to` | near-duplicate purpose / meaning | | 2 | `shares_purpose_with` | related goal, different mechanism | | 3 | `invalidates_with` | one makes the other wrong / stale | | 4 | `configured_by` | one is configured / parameterised by the other | | 5 | `depends_on` | one requires the other | | 6 | `contradicts` | the two make opposing claims | | 7 | `replaced_by` | one supersedes the other | ## 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 cross-encoder — with its single ranking logit swapped for an 8-class head, then **fine-tuned by sequence-level distillation** from a **Qwen2.5-7B-Instruct** teacher. The teacher read real documentation and labelled the relations between the entities it mentioned; those labels became the training targets. The pairs were normalised to the 8 classes, filtered for hallucinated / junk entities, windowed so both markers stay in context, and balanced with synthesised *no-relation* negatives and extra examples for the rarer classes. ## What's special - **A 7B's task in one small forward pass.** Relation extraction is normally done by prompting a large LLM; here it is a single ~sub-10 ms classification, cheap enough to run over a whole corpus. - **Joint (cross-encoder) reading.** The two entities and their context are read together in one pass, so the model can weigh how they actually relate — far more precise than comparing two independent embeddings. - **Abstain + confidence.** `NO_RELATION` plus a softmax threshold keep spurious pairs out of your graph. - **Ships as compiled engines** (TensorRT / OpenVINO, fp16) for production-speed inference, plus the ONNX for standalone use. ## Intended use Build a **knowledge graph**: for each pair of entities that co-occur in a passage, mark them and classify the relation. Wrap the two entities with `[E1]…[/E1]` and `[E2]…[/E2]`, tokenize as an *(empty-query, marked-text)* pair (the format it was trained on), run the engine, `argmax` the 8 logits, and drop `NO_RELATION` / low-confidence results. ```python import onnxruntime as ort, numpy as np from transformers import AutoTokenizer LABELS = ["NO_RELATION","semantically_similar_to","shares_purpose_with","invalidates_with", "configured_by","depends_on","contradicts","replaced_by"] tok = AutoTokenizer.from_pretrained(".") # bundled tokenizer incl. [E1]/[E2] markers sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"]) def classify(marked_text, max_len=256, tau=0.5): 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() i = int(p.argmax()) return (LABELS[i], float(p[i])) if i != 0 and p[i] >= tau else ("NO_RELATION", float(p[0])) # classify("The [E1]FAISS[/E1] vector index was replaced by the [E2]native IVF[/E2] backend.") # -> ('replaced_by', 0.7x) ``` ## What's in this repo Ready-to-run compiled engines, named per **runtime × GPU arch × OS** (single-profile — no length buckets): - **TensorRT** `code-daemon-relation-v1_{win_x64,linux_x64}_trt_sm_{86,89,120}.engine` — NVIDIA, fp16. - **OpenVINO** `code-daemon-relation-v1_ov_{cpu,igpu}_fp16_b16_s256.{xml,bin}` — Intel CPU / iGPU. - **Tokenizer** — `tokenizer.json` + `sentencepiece.bpe.model` + `tokenizer_config.json` (XLM-R SentencePiece with the 4 `[E1]/[E2]` marker tokens added). - **ONNX source** — `model.onnx` (+ `model.onnx.data`) FP32, for standalone `onnxruntime` / `optimum` use. > fp16 only: the mmarco format hits a known OpenVINO INT8 AccessViolation on the iGPU, so fp16 is shipped. ## Evaluation This is a **first distilled cut**, and the classes are naturally imbalanced (the teacher emits `semantically_similar_to` / `shares_purpose_with` far more often than `replaced_by` / `depends_on`). Reported honestly: **macro-F1 ≈ 0.33** on a held-out dev split, with per-class F1 in the **~0.2–0.44** band for the well-populated classes and lower on the thin tail classes, which the data under-samples and which are supplemented with synthetic examples. The abstain class plus a softmax threshold keep low-confidence pairs out of the graph. Metrics are advisory — for graph construction, spot-checking the edges it produces is the real test. ## License & training data Released under the **MIT license** (the mmarco base + XLM-R backbone are MIT/Apache; fine-tuned weights released MIT). | Source | Note | |---|---| | `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` (warm-start base) | **mMARCO ← MS MARCO → non-commercial research terms** | | distillation targets (Qwen2.5-7B-Instruct over open-source docs) | self-generated | | synthesised negatives + 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 is required. ## Attribution Warm-started from **[cross-encoder/mmarco-mMiniLMv2-L12-H384-v1](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1)**. Distilled from **Qwen2.5-7B-Instruct**. Backbone: XLM-RoBERTa.