tiny-graph-extractor-qwen3.5-0.8b-qlora

tiny-graph-extractor-qwen3.5-0.8b-lora

A LoRA adapter for unsloth/Qwen3.5-0.8B, trained with GRPO (on-policy RL with a reference-free reward) to extract entities and (head, relation, tail) triplets from text for knowledge graph ingestion.

This repository contains adapter weights only — the checkpoint from GRPO run 5, step 2000.

The goal of this project is to replace expensive frontier-LLM calls in the knowledge graph extraction pipeline of rst0070/knowledge-base with a tiny model that runs locally.

• Base model: unsloth/Qwen3.5-0.8B (Instruct), loaded 4-bit
• Method: QLoRA + GRPO — LoRA adapters on a 4-bit frozen base, trained with a reference-free reward (structure gate, dedup, grounding, and NLI-judged relation correctness/coverage using FactCG-DeBERTa as judge)
• Training prompts: sentences from REBEL (prompts only — the reward is reference-free, so REBEL's gold triplets are never used as targets)
• Hardware: single 16GB GPU (RTX 4060 Ti)
• Source: github.com/rst0070/tiny-entity-extractor

Results

Fixed evaluation reward over a 200-item test set (144 CrossRE + 30 curated + 26 real-news items), same reward function as training but with frozen weights/ramp so runs stay comparable. "Mean total incl. sentinels" counts unparseable outputs as −1.0 and is the quantity GRPO optimizes.

The adapter closes ~86% of the pretrained → Gemini gap on the optimized metric and matches Gemini on parse reliability. Note the reward has a ceiling below 1.0: the gold answers themselves score ~0.885 on NLI correctness, so a perfect extractor tops out around 0.93–0.94 (parsed-only) — Gemini sits on that ceiling.

Full per-component breakdowns and per-run analysis: data/result/README.md.

Known Contract

• Entities lean Wikipedia-style. Training prompts are REBEL (Wikipedia) sentences, so what counts as "an entity" follows Wikipedia conventions. Dates and numerical values are excluded by the prompt contract.
• Relations are free-form surface strings, grounded in the input text and judged by NLI — e.g. "was founded by", not Wikidata predicate labels. If your pipeline needs a fixed relation ontology, map these strings downstream.
• English, sentence/short-passage inputs. Evaluated on encyclopedic, multi-domain (CrossRE), and news text; long documents and other languages are untested.

Output Format

The model emits a single fenced JSON code block:

```json
{
  "entities": ["Entity A", "Entity B"],
  "relations": [
    {"head": "Entity A", "relation": "relation_type", "tail": "Entity B"}
  ]
}
```

Every entity referenced in a relation appears in the entities list.

Usage

Use the exact system prompt and user-message format the adapter was trained and evaluated with:

import re, json
import torch
from transformers import BitsAndBytesConfig, pipeline

SYSTEM_PROMPT = (
    "You are a knowledge graph extraction assistant.\n\n"
    "Given a text, extract entities and their relations. Output a single "
    "fenced JSON code block with this schema:\n"
    '- "entities": list of unique entity strings mentioned in the text\n'
    '- "relations": list of objects, each with "head", "relation", "tail" '
    "(all strings)\n\n"
    'Every entity referenced in a relation must appear in the "entities" list. '
    "Do not include dates or numerical values as entities. "
    "Output only the fenced JSON code block — no other text, preamble, or "
    "explanation."
)

text = (
    "Marie Curie was born in Warsaw and later moved to Paris, where she conducted her research at the University of Paris. She discovered polonium and radium."
)

# The adapter is a QLoRA checkpoint: it was trained and evaluated on a
# 4-bit NF4-quantized base, so load the base the same way.
pipe = pipeline(
    "text-generation",
    model="rst0070/tiny-graph-extractor-qwen3.5-0.8b-qlora",
    model_kwargs={
        "quantization_config": BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
        ),
    },
    max_new_tokens=1024,
)
messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": f"Extract entities and relations from the following text:\n\n{text}"},
]

response = pipe(messages)[0]["generated_text"][-1]["content"]
match = re.search(r"```(?:json)?\s*\n(.*?)\n```", response, re.DOTALL)
if match:
    print(json.loads(match.group(1)))

The BitsAndBytesConfig above matches the training/eval setup (Unsloth, load_in_4bit=True, i.e. bitsandbytes NF4). Loading the base without it (plain fp16/bf16) also works, but the reported numbers were measured with the 4-bit base, so quantization noise is part of them.

Training

GRPO — on-policy RL where the model learns directly from a reward on its own outputs:

1. Rollout: sample a group of completions per prompt from the current policy.
2. Score: compute a reference-free reward per completion — a structure gate (parseable JSON), dedup and grounding components, and NLI-judged relation correctness and coverage (FactCG-DeBERTa as the judge).
3. Update: normalize rewards within the group to advantages and take a clipped GRPO step, KL-regularized toward the base model.

This checkpoint is the result of 5 sequential runs: run 1 started GRPO from the pretrained model, and each later run continue-trained the previous run's best adapter on a fresh, disjoint 2000-sample REBEL slice, changing exactly one thing (reward weights, NLI acceptance ramp) so its effect is attributable. Progression on the optimized metric: 0.422 (pretrained) → 0.581 → 0.663 → 0.775 → 0.792 → 0.796 (run 5), vs. Gemini 2.5 Flash Lite at 0.858.

Limitations

• English only; other languages are untested.
• Free-form relation strings. No fixed relation ontology; downstream consumers needing canonical predicates (e.g. Wikidata) must map them.
• Sentence / short-passage inputs. Not validated on long documents.
• No factual grounding. The model extracts what it reads; it does not verify claims.
• Below-frontier quality. Matches Gemini 2.5 Flash Lite on parse reliability but still trails on relation correctness/coverage (0.796 vs 0.858 on the optimized metric).

License

Apache 2.0, inheriting from the base model and dataset license terms.