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

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. Distilled from Qwen2.5-7B-Instruct. Backbone: XLM-RoBERTa.