MIRA3_LEARNINGS.md

# Mira-3 Learnings — Building a Specialized Document→JSON Extraction SLM

**What this is.** A distilled, reusable playbook from building **Mira-3**: a ~4B, on-prem, fine-tuned small language model that reads a document and returns **structured, source-grounded JSON**. The beachhead is clinical extraction (lab reports, discharge summaries, intake forms → schema JSON), but **every lesson below generalizes to any document→JSON extraction task** — résumé/CV parsing, invoices, receipts, KYC forms, contracts, RFP/compliance docs. Where a lesson is clinical-specific, it's marked; most aren't.

Read this before starting a new extraction vertical (e.g. a "Mira-Résumé"). It will save you a training run and an eval that lies to you.

> **The one-line lesson:** *A fine-tuned extractor learns the exact surface form of its training data. If your training docs state facts one way and real docs state them another, the model silently fails on real docs — and an in-distribution synthetic eval will score ~1.0 and hide it. Measure on real, out-of-distribution documents, and measure recall, not just validity.*

---

## 0. The stack (what we actually shipped)

| Component | Choice | Why |
|---|---|---|
| Base model | `Qwen/Qwen3-4B-Instruct-2507` (Unsloth 4-bit: `unsloth/qwen3-4b-instruct-2507-unsloth-bnb-4bit`) | Strong instruction-following at a size that runs on a customer's single GPU. Open weights (no ToS trap). |
| Method | **QLoRA** (4-bit base + LoRA r16/α32), Unsloth, sequence packing | ~2× faster + low VRAM; fits a free Kaggle T4. |
| Trainer | TRL `SFTConfig`/`SFTTrainer`, full-sequence loss, packing on | Proven config; validity 1.0 on in-dist data. |
| Compute | **Kaggle free T4** (30h/wk), multi-session with lossless resume | Zero infra cost for v0. Paid GPU only for teacher generation. |
| Data (v0) | **Synthetic + permissively-licensed public only** | No DUA/PHI/copyright risk; retrain on real partner data at v1. |
| Output contract | JSON Schema with `x-grounding` / `x-match-on` extensions | Schema is a runtime input, not code — lets you onboard a new vertical without a code change. |

**Reusable assets in this repo** (paths are the concrete, working versions):
- Training kernel (resumable): `models/mira-3/notebooks/train_mira3_v11.py` + `RUNBOOK_v11.md`
- Eval kernel (2-phase, best-epoch): `models/mira-3/notebooks/eval_mira3_v11.py`
- Offline authoritative scorer: `scripts/mira3/build_v11_gate_scorecard.py`
- Exit gate: `scripts/mira3/phase25_gate_check.py`
- Grounding checker (schema-aware): `src/slm/data/_grounding.py`; field-F1: `src/slm/eval/_field_f1.py`

---

## 1. Training recipe that worked

- **QLoRA 4-bit + Unsloth + packing**, `max_seq_length` sized to cover ~100% of your token length distribution (we used 3072; a too-small 2048 silently truncated ~10% of rows in an earlier run and degraded quality). **Measure your length distribution first**; truncation is a silent quality leak.
- LoRA on all attention + MLP projections (`q,k,v,o,gate,up,down`), r=16, α=32, dropout 0.05, lr 2e-4 cosine, warmup 3%, effective batch 16 (2 × grad-accum 8), **3 epochs**.
- **Loss plateaus fast.** Loss went 1.36 → 0.15 within epoch 1 and only 0.15 → 0.13 across epochs 2–3. Most learning is epoch 1; later epochs are refinement. **Save a per-epoch adapter and pick the best by eval**, don't assume the last epoch is best (see §4).
- **Chat-template fidelity is non-negotiable.** Train with `apply_chat_template(..., add_generation_prompt=False)` over `[system,user,assistant]`; at inference use `add_generation_prompt=True`. Load the tokenizer **from the saved adapter dir** so the exact template travels with the model. A prompt/template mismatch between train and eval tanks validity for a reason that looks like a model failure.

---

## 2. Training under a hard wall-clock cap (Kaggle 12h) — lossless multi-session resume

A full 3-epoch run was ~20h of compute — impossible in one 12h session. The pattern that worked (generalizes to any pre-emptible/capped GPU: spot instances, Colab, CI runners):

1. **Wall-clock callback** forces a full `Trainer` checkpoint + graceful stop at ~10.5h (leave ≥1.5h headroom before the platform kill — the final save + output upload take real time; we lost a run's final adapter once by cutting it to 10.5h with too little margin).
2. **Resume = re-run the same kernel with the prior checkpoint mounted** as an input. The `Trainer` restores optimizer + LR schedule + RNG; completion is gated on `global_step >= max_steps` (**not** rounded epoch — a time-stop at epoch 2.99 is *not* "done").
3. **Freeze dependency versions to the exact set that wrote the checkpoint** before resuming. A lossless resume requires identical `transformers`/`peft`/`bitsandbytes`/`unsloth` — pin `==`, capture `env_versions.json` in session 1, and pin to those. Ranges drift and break resume.
4. **Archive-robust, fail-loud resume.** Upload the checkpoint as one tar with directory structure preserved (`--dir-mode skip`, **never** `--dir-mode zip` — it flattens `checkpoint-XXXX/` and silently restarts from step 0). If a checkpoint is *expected* (a checkpoint dataset is attached) but none is found, **crash** — don't silently burn a session training from scratch.

**Bug we hit (watch for it):** a per-epoch adapter saver keyed on `round(state.epoch)` mislabeled a mid-epoch budget-stop at epoch 1.57 as "epoch_2". Guard epoch-boundary saves with `abs(epoch - round(epoch)) < 0.05` so a fractional stop doesn't write a bogus artifact.

---

## 3. Infra/Kaggle gotchas (each cost us time)

- **Inline `pip install` at the top of CLI-pushed kernels** — they don't inherit the notebook UI's packages. One pinned command, fail fast (before GPU time).
- **Split training and eval into separate notebooks.** Never run both in one session (12h timeout). Eval is a separate, shorter GPU session.
- **`kaggle kernels output` is flaky** — it returns a *different partial subset of files per call*, and the log often needs a second pull. Re-run the pull 2–3× and verify the critical files (metrics, adapters, log) are non-zero before trusting them. A "0-byte final adapter" locally was a download artifact, not a training failure.
- **Batch (CLI-pushed) kernels don't stream logs.** The whole log publishes only on commit. "No new logs for hours" during a RUNNING batch kernel is normal, not a stall. Status (`RUNNING`/`COMPLETE`) is your only live signal; poll it.
- **Read adapters across kernels via `kernel_sources`**, not by re-uploading — the eval kernel mounts the training kernel's `/kaggle/working/output` directly.
- The training output tree may contain a `_resume/` staging dir with **stale copies** of adapters. When discovering adapters, exclude `_resume/` and `checkpoints/` and require `adapter_config.json`, or you may silently score a stale/partial adapter.

---

## 4. Eval methodology — the part that actually determines whether you have a product

This is where most of the value is. **Evaluation is harder and more important than training** for extraction models.

### 4a. Mechanics
- **Best-epoch selection:** score every per-epoch adapter and pick the winner. Load each with the *same* loader the training used (for Unsloth 4-bit adapters, `FastLanguageModel`, **not** vanilla `PeftModel` — quantization mismatch produces garbage).
- **Two-phase, quota-frugal eval:** Phase A scores all epochs on the labeled set (cheap, decides best epoch first so a timeout still yields it); Phase B runs the full probe suite on the winner only.
- **In-kernel metrics are for monitoring/ranking; authoritative metrics are computed offline** on the saved raw predictions (schema-aware grounding, macro-F1) with the real library. Save raw predictions in the exact schema your offline loader expects (`extra="forbid"` loaders reject stray keys — we crashed the offline scorer by writing a `gold` key and omitting `gen_time`/`pred_len`).
- **Adversarially review the eval/train kernel *before* pushing it.** A multi-agent review (independent reviewers per failure-mode: mounts, adapter-loading, prompt fidelity, metric logic, output-format-vs-offline-loader, time-budget) caught 3 real bugs that would each have wasted a ~5h GPU session. Cheap insurance for anything that costs a GPU run.

### 4b. Model-specific output gotcha — the `<think>` prefix
Qwen3-Instruct emits a leading `<think>...</think>` block (often empty) before the JSON. Our raw validity read **0%** until we stripped it — the JSON was perfect underneath. **Strip the reasoning wrapper before parsing** (and the serving layer must too). Check your base model's default output framing before trusting a "0% valid" result.

### 4c. THE BIG ONE — in-distribution synthetic eval lies

Our labeled `test_gold` (synthetic, same generator as training) scored **field-F1 = 1.000, validity = 1.000, 200/200 predictions byte-identical to gold.** This was **not** leakage (0/800 exact overlap, 4.4% template overlap) and the metric was **not** broken (corrupting a field dropped F1 to 0.92). It was real — and **meaningless as a generalization signal.** The synthetic docs have a near-deterministic source→label mapping; the model learned it perfectly. The *previous* model scored 1.0 on it too.

**If your held-out set comes from the same generator/distribution as training, a high score tells you the model fit the generator — nothing about the real world.**

### 4d. Build a labeled REAL-doc benchmark (how)
Real-doc probes (actual clinical transcripts) were **unlabeled**, so we could only measure validity + hallucination, not F1. To get an honest field-F1 we built a small labeled set:
1. Select a **diverse** sample (we took 24 real docs, one per specialty/domain).
2. **Double-annotate + adjudicate:** two independent annotators produce gold from *source + schema only* (never shown the model's output), then an adjudicator resolves disagreements strictly against the text. (We used LLM annotators + adjudication — a defensible *directional* signal; replace with human gold for certification.)
3. Score the model's saved predictions against the adjudicated gold with schema-aware field-F1 (exclude meta fields like `extraction_notes`; coerce malformed list items).

### 4e. What the honest eval revealed (the pattern to expect)
Real-doc **macro-F1 = 0.61** (vs synthetic 1.0). The shape matters more than the number:

- **High precision (0.89–0.98) on every field, but low recall.** The model is *safe* (what it emits is right) but *misses content*.
- **The blind spot:** it extracted **~0 diagnoses from real narrative notes** (docs with 6–10 explicit diagnoses → output `[]`), and under-extracted procedures (recall 0.33). It nailed list-structured fields (medications, labs, vitals, allergies, encounter date/dept).

**Root cause = data, not model.** The synthetic training data only ever presents diagnoses/procedures in **structured `Diagnoses:` blocks with bracketed codes**. Real notes state them in **prose** ("patient has DM type 2, HTN…"). The model learned "extract diagnoses from the Diagnoses: section" and never learned to read narrative. **This is the single most important, most generalizable lesson.**

---

## 5. Data lessons (all generalize)

1. **Structured-vs-narrative coverage is everything.** Whatever field a downstream user cares about, your training data must present it in **every surface form real documents use** — structured lists *and* inline prose, tables *and* sentences. If a form is only ever seen one way in training, the model fails on the other. *This is the recall killer.* (For résumés: a "skill" appears as a bulleted list, a comma-separated line, and buried in a job-description paragraph — train on all three.)
2. **Preserve source formatting; don't canonicalize in labels.** Our training labels rendered a lab unit one fixed way (`x10^9/L`) with zero variety, so the model "corrected" the real docs' `10^9/L` to `x10^9/L` — violating "preserve exactly." **Fix: render source values/units in varied equivalent forms and have the label copy the source verbatim.** Teach copying, not canonicalization.
3. **Give the verifier principled equivalences.** A correct-unit reformat (`10^9/L ≡ x10^9/L`) is not a hallucination and shouldn't route to a human. Add narrow, meaning-preserving equivalences to the grounding checker (scoped to the right field, so it can never ground a fabricated *value*). Verified precision: it must not "fix" a genuine fabrication (e.g., inventing `Glucose (Fasting)` from `glucose`).
4. **Audit your generator empirically.** Run the grounding checker *on your own training labels*: any label value not grounded in its own source is a canonicalization you're teaching. Ours was clean (0.02%) — which is how we learned the unit issue was a *diversity* gap, not a labeling bug. This 5-minute check reframes fixes.

---

## 6. Metrics & acceptance framework (use for any vertical)

| Metric | How | Bar | Notes |
|---|---|---|---|
| JSON validity | parse to schema-valid object (after stripping any reasoning wrapper) | ≥99% raw, 100% constrained | 100% is achievable at serving via constrained decoding; measure raw to see model behavior. |
| Field-F1 (**precision + recall separately**) | schema-aware match (`x-match-on` keys); exact for scalars | ≥0.90 **on real docs** | Report P and R separately — a low-recall/high-precision model looks fine on F1-averaged-with-validity but misses content. |
| Hallucination | grounding: every extracted value must appear in source (schema `x-grounding: exempt` for classifications) | 0 (post-verifier) | Raw model is never 0 on real docs; the *system* hits 0 via verifier + human-in-loop. |
| Identifier leak | regex for names/MRN/DOB/SSN/phone/email | 0 | Schema stores no identifiers by design. |
| Beats baseline | vs the prior model AND base zero-shot | strictly better | Base zero-shot ≈ 0% valid JSON; the meaningful bar is the prior fine-tune. |

**Do not report a single blended number.** Report: validity, **per-field precision AND recall**, hallucination, leak — on **real** docs. The blended F1 on synthetic data is the metric that lies.

---

## 7. Porting to a new vertical (e.g. Résumé / CV parsing) — checklist

The pipeline is schema-driven, so a new vertical is mostly data + schema:

1. **Schema.** Define the JSON contract with `x-grounding` (which fields must be verbatim-grounded vs. are classifications/derived) and `x-match-on` (list-item match keys). E.g. résumé: `experience[]` matched on `(company, title)`, `education[]` on `(institution, degree)`, `skills[]` on `name`. Mark `summary`/`document_type` as grounding-exempt.
2. **Data — cover every surface form (see §5.1).** For each field, generate/collect docs where it appears structured *and* in prose. This is where you prevent the recall blind spot before it happens.
3. **Preserve formatting (§5.2).** Dates, phone numbers, degree abbreviations (`BS` vs `Bachelor of Science`) — vary them in sources and copy verbatim in labels; add equivalences to the verifier, not the training labels.
4. **Train** with the resumable QLoRA recipe (§1–2); save per-epoch adapters.
5. **Eval honestly (§4).** Build a small **real** résumé benchmark with double-annotation + adjudication. Measure per-field precision *and* recall. Expect (and hunt for) the high-precision/low-recall blind spot on the fields real documents phrase in prose.
6. **Gate + verifier + human-in-loop** identical to clinical — the whole workflow (extract → ground → verify → correct → audit) is the product, not the single model call.

---

## 8. TL;DR — the five things to remember

1. **The model learns the surface form of your data.** Train on every way real documents phrase a field, especially **narrative prose**, or you get high precision and terrible recall on real docs.
2. **In-distribution synthetic eval scores ~1.0 and means nothing.** Build a labeled **real-doc** benchmark; measure **recall**, not just validity/F1.
3. **Preserve source formatting in labels; put equivalences in the verifier.** Don't teach the model to canonicalize.
4. **Adversarially review any kernel before spending a GPU session;** strip the base model's reasoning wrapper before parsing; compute authoritative metrics offline.
5. **For capped GPUs, do lossless multi-session resume** with frozen deps, step-gated completion, and fail-loud archive handling.

*Written from the Mira-3 v11 build (clinical extraction, Qwen3-4B QLoRA). v11 = valid JSON 100%, identifier-leak 0, beats the prior model — but real-doc field-F1 0.61 with a diagnoses-recall blind spot that the synthetic eval hid. The blind spot is a data problem, and it's the v1 priority.*