mira-3-v11 / MIRA3_LEARNINGS.md

Mira-3 v11: adapter + per-epoch + training metrics + eval results + learnings

99da1a4 verified about 2 months ago

16.7 kB

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):

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).
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").
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.
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:

Select a diverse sample (we took 24 real docs, one per specialty/domain).
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.)
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)

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.)
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.
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).
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:

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.
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.
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.
Train with the resumable QLoRA recipe (§1–2); save per-epoch adapters.
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.
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

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.
In-distribution synthetic eval scores ~1.0 and means nothing. Build a labeled real-doc benchmark; measure recall, not just validity/F1.
Preserve source formatting in labels; put equivalences in the verifier. Don't teach the model to canonicalize.
Adversarially review any kernel before spending a GPU session; strip the base model's reasoning wrapper before parsing; compute authoritative metrics offline.
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.