runtime/jevlike/adapt.py

"""Adapt a checkpoint to your own labeled task: fine-tune (with replay) + calibrate + before/after report.

    python -m jevlike.adapt --data labels.csv --spec task.json --base checkpoints/jevlike-large-nf --out checkpoints/my-task
    python -m jevlike.adapt --records my_records.jsonl --base checkpoints/jevlike-large-nf --out checkpoints/my-task
    python -m jevlike.adapt --data labels.csv --spec task.json --out checkpoints/my-task --calibrate-only   # < ~100 rows

Task spec (JSON or YAML): questions in the public API shape plus where each one's label lives.

    {"text_column": "headline",                      # or a list of columns -> JSON-object state
     "questions": {
       "topic": {"type": "choice", "instructions": "What is this news item mainly about?",
                 "criteria": {"markets": "stocks, bonds, FX", "earnings": "company results, guidance"},
                 "label_column": "topic", "label_map": {"Mkts": "markets"}},
       "tone":  {"type": "score", "instructions": "How positive is the news for the company?",
                 "criteria": ["negative", "neutral", "positive"]},
       "guidance": {"type": "noul", "instructions": "The company changes its financial guidance"}}}

The "questions" wrapper is optional (a bare {name: question} dict works). A question's label column
is, in order: --label-col NAME=COL, its "label_column", the spec's top-level "label_columns"
{name: col}, a column named like the question, or "label" when there is a single question.
Label values are matched to option keys (then case-insensitively), score level texts, an integer
option index, or true/false/yes/no/1/0 for noul, after "label_map". Empty cells are skipped.

Pipeline (writes everything under --out; the base checkpoint is never modified):
  1. rows -> Records; stratified train/val/test split (default 60/20/20) grouped by text
  2. base checkpoint evaluated zero-shot on test, and with temperatures re-fitted on your val
  3. fine-tune from the base (python -m jevlike.train --init-from) on your train split repeated
     --epochs times plus a replay sample of general data, best checkpoint by val NLL
  4. calibrate on your val (python -m jevlike.calibrate); cells with too little data keep a
     general-data calibration; evaluate on test (python -m jevlike.evaluate) and on general data
  5. OUT/adaptation_report.{md,json}: accuracy / NLL / ECE / per-class accuracy before and after
--calibrate-only skips 3: OUT becomes a copy of the base with a calibration fitted on your data.
"""
from __future__ import annotations

import argparse
import csv
import datetime as _dt
import json
import math
import os
import random
import re
import shlex
import shutil
import subprocess
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional, Union

from jevlike.types import Choice, Noul, Question, Record, Score, read_jsonl, write_jsonl

REPO = Path(__file__).resolve().parents[1]
DEFAULT_BASE = "checkpoints/jevlike-large-nf"
DEFAULT_REPLAY = ("data/mix3/train.jsonl", "data/train.jsonl")
DEFAULT_GENERAL = ("data/mix3/val.jsonl", "data/val.jsonl")
TRUE = {"1", "true", "yes", "y", "t"}
FALSE = {"0", "false", "no", "n", "f"}
_Q_FIELDS = {"type", "instructions", "criteria", "label_column", "column", "label_map"}
_META = {"name", "description", "text_column", "id_column", "label_columns", "questions"}


# ---------------------------------------------------------------- task spec

@dataclass
class QSpec:
    name: str
    question: Question
    label_column: Optional[str] = None
    label_map: dict[str, str] = field(default_factory=dict)
    level_keys: Optional[list[str]] = None  # score given as {key: description}


@dataclass
class TaskSpec:
    questions: list[QSpec]
    text_column: Union[str, list[str]] = "text"
    id_column: Optional[str] = None
    name: str = "task"


def load_spec_file(path: str) -> dict:
    text = Path(path).read_text(encoding="utf-8")
    if Path(path).suffix.lower() in (".yaml", ".yml"):
        import yaml
        return yaml.safe_load(text)
    return json.loads(text)


def _question(name: str, d: dict) -> QSpec:
    if not isinstance(d, dict):
        raise ValueError(f"question {name!r}: expected an object, got {type(d).__name__}")
    extra = set(d) - _Q_FIELDS
    if extra:
        raise ValueError(f"question {name!r}: unknown field(s) {sorted(extra)}; allowed {sorted(_Q_FIELDS)}")
    t, instr, crit = d.get("type"), d.get("instructions"), d.get("criteria")
    if not isinstance(instr, str) or not instr.strip():
        raise ValueError(f"question {name!r}: 'instructions' must be a non-empty string")
    level_keys = None
    if t == "choice":
        if not isinstance(crit, (dict, list)):
            raise ValueError(f"question {name!r}: choice needs 'criteria' as {{key: description}} or [key, ...]")
        q = Question.from_public(Choice(instr, {str(k): str(v or "") for k, v in crit.items()}
                                        if isinstance(crit, dict) else [str(k) for k in crit]))
    elif t == "score":
        if isinstance(crit, dict):  # same convention as the server: the model reads "key: description"
            level_keys = [str(k) for k in crit]
            crit = [f"{k}: {v}" if v else str(k) for k, v in crit.items()]
        if not isinstance(crit, list):
            raise ValueError(f"question {name!r}: score needs 'criteria' as a list of levels (lowest first)")
        q = Question.from_public(Score(instr, [str(c) for c in crit]))
    elif t == "noul":
        if crit:
            raise ValueError(f"question {name!r}: noul takes no criteria; put the statement in 'instructions'")
        q = Question.from_public(Noul(instr))
    else:
        raise ValueError(f"question {name!r}: type must be choice, score or noul (got {t!r})")
    try:
        q.validate()
    except AssertionError as e:
        raise ValueError(f"question {name!r}: {e}") from None
    lm = d.get("label_map") or {}
    if not isinstance(lm, dict):
        raise ValueError(f"question {name!r}: label_map must be an object {{file value: key}}")
    return QSpec(name, q, d.get("label_column") or d.get("column"), {str(k): str(v) for k, v in lm.items()},
                 level_keys)


def parse_spec(d: dict, text_column: Optional[str] = None, id_column: Optional[str] = None,
               label_columns: Optional[dict[str, str]] = None) -> TaskSpec:
    """Spec dict (see module doc) -> TaskSpec. CLI overrides win over the file."""
    if not isinstance(d, dict):
        raise ValueError("task spec must be a JSON/YAML object")
    if isinstance(d.get("questions"), dict):
        qd = d["questions"]
        extra = set(d) - _META
    else:  # bare {name: question}
        qd = {k: v for k, v in d.items() if k not in _META}
        extra = {k for k, v in qd.items() if not (isinstance(v, dict) and "type" in v)}
    if extra:
        raise ValueError(f"task spec: unknown top-level key(s) {sorted(extra)}")
    if not qd:
        raise ValueError("task spec has no questions")
    qs = [_question(str(k), v) for k, v in qd.items()]
    top_map = {str(k): str(v) for k, v in (d.get("label_columns") or {}).items()}
    for q in qs:
        q.label_column = (label_columns or {}).get(q.name) or q.label_column or top_map.get(q.name)
    unknown = set(label_columns or {}) - {q.name for q in qs}
    if unknown:
        raise ValueError(f"--label-col for unknown question(s) {sorted(unknown)}")
    return TaskSpec(qs, text_column or d.get("text_column") or "text", id_column or d.get("id_column"),
                    str(d.get("name") or "task"))


# ---------------------------------------------------------------- labeled table -> Records

def load_table(path: str) -> list[dict]:
    p = Path(path)
    suf = p.suffix.lower()
    if suf in (".jsonl", ".ndjson"):
        with open(p, encoding="utf-8") as f:
            return [json.loads(line) for line in f if line.strip()]
    if suf == ".json":
        rows = json.loads(p.read_text(encoding="utf-8"))
        if not (isinstance(rows, list) and all(isinstance(r, dict) for r in rows)):
            raise ValueError(f"{path}: expected a JSON array of objects")
        return rows
    if suf in (".csv", ".tsv", ".txt"):
        csv.field_size_limit(min(sys.maxsize, 2**31 - 1))
        with open(p, encoding="utf-8-sig", newline="") as f:  # -sig: drop a BOM
            return list(csv.DictReader(f, delimiter="\t" if suf == ".tsv" else ","))
    raise ValueError(f"{path}: unsupported file type {suf!r} (use .csv, .tsv, .jsonl or .json)")


def _as_str(v: Any) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, float) and v.is_integer():
        return str(int(v))
    return "" if v is None else str(v).strip()


def _int_like(s: str) -> Optional[int]:
    try:
        f = float(s)
    except ValueError:
        return None
    return int(f) if f.is_integer() else None


def resolve_label(qs: QSpec, value: Any) -> Optional[int]:
    """File cell -> label index; None for an empty cell; ValueError if it matches nothing."""
    v = _as_str(value)
    if v == "":
        return None
    if qs.label_map:
        lm_ci = {k.lower(): t for k, t in qs.label_map.items()}
        v = qs.label_map.get(v, lm_ci.get(v.lower(), v))
    q = qs.question
    if q.type == "noul":
        if v.lower() in TRUE:
            return 1
        if v.lower() in FALSE:
            return 0
        raise ValueError(v)
    names = [qs.level_keys or [], [o.key for o in q.options]] if q.type == "choice" else \
        [qs.level_keys or [], [o.text for o in q.options]]
    for cand in names:
        if v in cand:
            return cand.index(v)
    for cand in names:
        low = [c.lower() for c in cand]
        if v.lower() in low:
            return low.index(v.lower())
    i = _int_like(v)
    if i is not None and 0 <= i < len(q.options):
        return i
    raise ValueError(v)


def _state_of(row: dict, text_column: Union[str, list[str]]) -> str:
    if isinstance(text_column, list):  # several columns -> JSON object state, as the server sends it
        return json.dumps({c: row.get(c) for c in text_column}, indent=2, ensure_ascii=False)
    v = row.get(text_column)
    if isinstance(v, (dict, list)):
        return json.dumps(v, indent=2, ensure_ascii=False)
    return "" if v is None else str(v).strip()


def rows_to_records(rows: list[dict], spec: TaskSpec, skip_bad: bool = False) -> tuple[list[Record], dict]:
    """Returns (records, info) with info = resolved label columns and skipped counts."""
    if not rows:
        raise ValueError("the labeled file has no rows")
    cols = list(dict.fromkeys(k for r in rows[:1000] for k in r))
    tcols = spec.text_column if isinstance(spec.text_column, list) else [spec.text_column]
    missing = [c for c in tcols if c not in cols]
    if missing:
        raise ValueError(f"text column(s) {missing} not in the file; columns are {cols}")
    label_cols = {}
    for q in spec.questions:
        c = q.label_column or (q.name if q.name in cols else "label" if len(spec.questions) == 1 else q.name)
        if c not in cols:
            raise ValueError(f"question {q.name!r}: label column {c!r} not in the file; columns are {cols} "
                             f"(set label_column in the spec or --label-col {q.name}=COLUMN)")
        label_cols[q.name] = c
    bad: dict[str, Counter] = defaultdict(Counter)
    recs, skipped = [], Counter()
    for i, row in enumerate(rows):
        state = _state_of(row, spec.text_column)
        if not state.strip():
            skipped["empty text"] += 1
            continue
        rid = (_as_str(row.get(spec.id_column)) if spec.id_column else "") or f"row{i}"
        for q in spec.questions:
            try:
                y = resolve_label(q, row.get(label_cols[q.name]))
            except ValueError:
                bad[q.name][_as_str(row.get(label_cols[q.name]))] += 1
                continue
            if y is None:
                skipped[f"{q.name}: empty label"] += 1
                continue
            r = Record(f"{rid}-{q.name}", q.name, state, q.question, y)
            r.validate()
            recs.append(r)
    for qn, c in bad.items():
        skipped[f"{qn}: unmatched label"] += sum(c.values())
    if bad and not skip_bad:
        lines = []
        for q in spec.questions:
            if q.name in bad:
                allowed = ["true/false"] if q.question.type == "noul" else \
                    (q.level_keys or [o.key if q.question.type == "choice" else o.text for o in q.question.options])
                lines.append(f"  {q.name}: {dict(bad[q.name].most_common(10))} (allowed: {allowed}, or an index)")
        raise ValueError("label values that match no option (add a label_map, or pass --skip-bad-labels):\n"
                         + "\n".join(lines))
    if not recs:
        raise ValueError("no usable labeled examples")
    return recs, {"label_columns": label_cols, "skipped": dict(skipped), "rows": len(rows)}


# ---------------------------------------------------------------- split

def _norm(s: str) -> str:
    return " ".join(s.split()).lower()


def parse_split(s: str) -> tuple[float, float, float]:
    parts = [float(x) for x in re.split(r"[,/ ]+", s.strip()) if x]
    if len(parts) != 3 or min(parts) < 0 or sum(parts) <= 0:
        raise ValueError(f"--split wants three non-negative numbers train,val,test (got {s!r})")
    tot = sum(parts)
    return parts[0] / tot, parts[1] / tot, parts[2] / tot


def split_records(records: list[Record], fractions=(0.6, 0.2, 0.2), seed: int = 13
                  ) -> tuple[list[Record], list[Record], list[Record]]:
    """Stratified (by the first question's label of each text) split, grouped by normalized text so
    no text lands in two splits. Deterministic for a seed; global proportions within +-1 group."""
    groups: dict[str, list[Record]] = {}
    for r in records:
        groups.setdefault(_norm(r.state), []).append(r)
    strata: dict[tuple, list[str]] = defaultdict(list)
    for k, rs in groups.items():
        strata[(rs[0].source, rs[0].label)].append(k)
    rng = random.Random(seed)
    order: list[str] = []
    for s in sorted(strata, key=str):
        ks = strata[s]
        rng.shuffle(ks)
        order += ks
    # systematic allocation: each group goes to the split furthest below its quota, so every
    # stratum (contiguous in `order`) is split in proportion
    counts, out = [0, 0, 0], ([], [], [])
    for j, k in enumerate(order):
        s = max(range(3), key=lambda i: (fractions[i] * (j + 1) - counts[i], -i))
        counts[s] += 1
        out[s].extend(groups[k])
    return out


def class_name(q: Question, label: int) -> str:
    if q.type == "noul":
        return ("false", "true")[label]
    o = q.options[label]
    return o.key if q.type == "choice" else (f"{o.key}: {o.text}" if o.text else o.key)


def class_counts(records: list[Record]) -> dict[str, dict[str, int]]:
    out: dict[str, Counter] = defaultdict(Counter)
    for r in records:
        out[r.source][class_name(r.question, r.label)] += 1
    return {s: dict(sorted(c.items())) for s, c in sorted(out.items())}


def write_by_class(records: list[Record], d: Path) -> dict[str, tuple[str, str]]:
    """One JSONL per (question, gold class) so evaluate's by_file table gives per-class accuracy."""
    d.mkdir(parents=True, exist_ok=True)
    by: dict[tuple[str, str], list[Record]] = defaultdict(list)
    for r in records:
        by[(r.source, class_name(r.question, r.label))].append(r)
    files: dict[str, tuple[str, str]] = {}
    for (src, cls), rs in sorted(by.items()):
        stem = re.sub(r"[^A-Za-z0-9._-]+", "_", f"{src}__{cls}")[:80]
        name, n = f"{stem}.jsonl", 1
        while name in files:
            n += 1
            name = f"{stem}_{n}.jsonl"
        write_jsonl(str(d / name), rs)
        files[name] = (src, cls)
    return files


# ---------------------------------------------------------------- replay / general samples

def sample_jsonl(path: str, k: int, seed: int, exclude_states: set[str] = frozenset()) -> list[Record]:
    """Uniform sample of k records (reservoir over lines), skipping states in exclude_states."""
    if k <= 0:
        return []
    rng, keep, n = random.Random(seed), [], 0
    want = k + max(20, k // 10)  # margin for excluded states
    with open(path, encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            if len(keep) < want:
                keep.append(line)
            else:
                j = rng.randrange(n + 1)
                if j < want:
                    keep[j] = line
            n += 1
    rng.shuffle(keep)
    recs = [r for r in (Record.from_json(x) for x in keep) if _norm(r.state) not in exclude_states]
    return recs[:k]


def first_existing(paths) -> Optional[str]:
    """First path that exists, relative to the cwd or else to the repo root."""
    for p in paths:
        for c in (Path(p), REPO / p):
            if p and c.exists():
                return str(c.resolve())
    return None


def estimate_tokens(r: Record) -> float:
    q = r.question
    chars = len(r.state) + len(q.instructions) + sum(len(o.key) + len(o.text) + 3 for o in q.options)
    return chars / 4.5 + 6 + 2 * q.n_options  # ~ModernBERT tokens (checked on mix3 + FOMC)


def plan_steps(records: list[Record], max_tokens: int, max_batch: int, grad_accum: int, max_len: int) -> int:
    """Optimizer steps for about one pass over `records` with token-budget micro-batches."""
    toks = sum(min(max_len, estimate_tokens(r)) for r in records)
    n_micro = max(toks / (0.92 * max_tokens), len(records) / max_batch)  # length-bucketed packing
    return max(10, math.ceil(n_micro / grad_accum))


# ---------------------------------------------------------------- calibration merge

def merge_calibration(user: dict, fallback: Optional[dict], min_count: int) -> dict:
    """Keep the user-fitted temperature where the user's val had >= min_count questions of that
    (type, option bucket) - or of that type, for a bucket it has any of - else the fallback's."""
    out = json.loads(json.dumps(user))
    cells: dict[str, dict[str, str]] = {}
    for t, buckets in user["temperatures"].items():
        cnt = user["counts"][t]
        type_n = sum(cnt.values())
        cells[t] = {}
        for b in buckets:
            if cnt.get(b, 0) >= min_count:
                cells[t][b] = f"yours (n={cnt[b]})"
            elif cnt.get(b, 0) > 0 and type_n >= min_count:
                cells[t][b] = f"yours, per-type (n={type_n})"
            elif fallback:
                out["temperatures"][t][b] = fallback["temperatures"].get(t, {}).get(b, 1.0)
                cells[t][b] = "fallback"
            else:
                cells[t][b] = "T=1 (no data)"
        if type_n < min_count and fallback:
            out["type_temperature"][t] = fallback.get("type_temperature", {}).get(t, 1.0)
    n_all = sum(sum(c.values()) for c in user["counts"].values())
    esc = "yours"
    if n_all < min_count and fallback:
        out["escalate"], esc = fallback["escalate"], "fallback"
    out["adapt_merge"] = {"min_count": min_count, "cells": cells, "escalate": esc}
    return out


# ---------------------------------------------------------------- subprocess plumbing

_HELP_CACHE: dict[str, set[str]] = {}


def cli_flags(module: str) -> set[str]:
    """The --flags a jevlike CLI accepts (parsed from --help), so optional pass-throughs are safe."""
    if module not in _HELP_CACHE:
        p = subprocess.run([sys.executable, "-m", module, "--help"], cwd=REPO, capture_output=True, text=True)
        _HELP_CACHE[module] = set(re.findall(r"--[a-z][a-z0-9-]*", p.stdout))
    return _HELP_CACHE[module]


def run_cmd(cmd: list[str], log_path: Path, env: Optional[dict] = None) -> None:
    line = "$ " + shlex.join(cmd)
    print(line, flush=True)
    tail: list[str] = []
    with open(log_path, "a", encoding="utf-8") as lf:
        lf.write(line + "\n")
        p = subprocess.Popen(cmd, cwd=REPO, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                             text=True, bufsize=1)
        for out in p.stdout:
            lf.write(out)
            print("  | " + out, end="", flush=True)
            tail = (tail + [out])[-30:]
        rc = p.wait()
    if rc:
        raise RuntimeError(f"command failed with exit code {rc} (log: {log_path}):\n{line}\n" + "".join(tail))


# ---------------------------------------------------------------- report

def metrics_of(rep: dict) -> dict:
    m = rep["overall"]["model"]
    keys = ("n", "accuracy", "nll", "brier", "ece", "accuracy_raw", "nll_raw", "ece_raw")
    out = {k: m.get(k) for k in keys}
    out["escalate_auroc"] = (rep.get("escalate") or {}).get("auroc_error")
    out["majority_accuracy"] = rep["overall"]["prior"].get("accuracy")
    return out


def _fmt(x, pct=False) -> str:
    if x is None:
        return "-"
    return f"{100 * x:.1f}%" if pct else f"{x:.4f}"


def _ci(p, n) -> str:
    return f" ±{196 * math.sqrt(max(p * (1 - p), 1e-9) / n):.1f}" if p is not None and n else ""


SYSTEM_LABELS = {"base": "base, zero-shot (shipped calibration)", "base_recal": "base + calibration fitted on your val",
                 "adapted": "adapted (fine-tuned + calibrated on your val)"}


def render_report(r: dict) -> str:
    L = [f"# Adaptation report: {Path(r['out']).name}", "",
         f"- mode: **{r['mode']}**; base `{r['base']}` -> output `{r['out']}`; created {r['created']}, "
         f"took {r['elapsed_s'] / 60:.1f} min",
         f"- data: {r['data']['source']}: {r['data']['n_records']} labeled questions over "
         f"{r['data']['n_texts']} texts; split train/val/test = " +
         " / ".join(f"{r['data']['split'][s]}" for s in ("train", "val", "test")) + " questions (grouped by text)"]
    if r["data"].get("skipped"):
        L.append(f"- skipped: {r['data']['skipped']}")
    L += ["", "## Test set, before and after", "",
          "| system | acc (95% CI) | NLL | Brier | ECE | ECE raw | escalate AUROC |", "|---|---|---|---|---|---|---|"]
    for k, m in r["systems"].items():
        L.append(f"| {SYSTEM_LABELS.get(k, k)} | {_fmt(m['accuracy'], True)}{_ci(m['accuracy'], m['n'])} | "
                 f"{_fmt(m['nll'])} | {_fmt(m['brier'])} | {_fmt(m['ece'])} | {_fmt(m['ece_raw'])} | "
                 f"{_fmt(m['escalate_auroc'])} |")
    anym = next(iter(r["systems"].values()))
    L.append(f"| majority class (fit on train) | {_fmt(anym['majority_accuracy'], True)} | | | | | |")
    L += ["", "Recalibration changes NLL/ECE, never accuracy (temperature scaling keeps the argmax).", ""]
    acc_sys = [s for s in r["systems"] if s != "base_recal"]
    L += ["## Per question", "", "| question | n | " + " | ".join(f"{s} acc | {s} balanced acc" for s in acc_sys) + " |",
          "|---|---|" + "---|---|" * len(acc_sys)]
    for qn, d in r["per_question"].items():
        L.append(f"| {qn} | {d['n']} | " + " | ".join(
            f"{_fmt(d[s]['accuracy'], True)} | {_fmt(d[s]['balanced_accuracy'], True)}" for s in acc_sys) + " |")
    L += ["", "## Per class (recall: accuracy on test questions whose gold answer is that class)", "",
          "| question | class | n | " + " | ".join(acc_sys) + " |", "|---|---|---|" + "---|" * len(acc_sys)]
    for qn, classes in r["per_class"].items():
        for cls, d in classes.items():
            L.append(f"| {qn} | {cls} | {d['n']} | " + " | ".join(_fmt(d.get(s), True) for s in acc_sys) + " |")
    if r.get("general"):
        g = r["general"]
        L += ["", f"## General data (forgetting check: {g['n']} held-out questions from `{g['file']}`)", "",
              "| system | acc | NLL | ECE | NLL raw | ECE raw |", "|---|---|---|---|---|---|"]
        for s in ("base", "adapted"):
            m = g.get(s)
            if m:
                L.append(f"| {s} | {_fmt(m['accuracy'], True)} | {_fmt(m['nll'])} | {_fmt(m['ece'])} | "
                         f"{_fmt(m['nll_raw'])} | {_fmt(m['ece_raw'])} |")
        L.append("\nThe adapted model's general-data temperatures come from a separate general sample.")
    if r.get("training"):
        t = r["training"]
        L += ["", "## Training", "",
              f"- {t['n_user_train']} of your train questions x {t['epochs']} passes + {t['n_replay']} replay "
              f"questions from `{t['replay_file']}` = {t['n_mix']} questions, {t['max_steps']} optimizer steps, "
              f"eval every {t['eval_every']} steps, best checkpoint by val NLL",
              f"- command: `{t['command']}`"]
    c = r["calibration"]
    L += ["", "## Calibration", "",
          f"- fitted on your val split ({r['data']['split']['val']} questions); a (type, option-count) cell needs "
          f">= {c['min_count']} val questions, else it keeps the fallback ({c['fallback']})",
          f"- cells: {json.dumps({t: {b: s for b, s in d.items() if s != 'fallback'} for t, d in c['cells'].items()})}",
          "- temperatures are shared by every question of the same type and option count, including ones "
          "outside your task"]
    if r.get("notes"):
        L += ["", "## Notes", ""] + [f"- {n}" for n in r["notes"]]
    L += ["", "## Files", ""] + [f"- {k}: `{v}`" for k, v in r["files"].items()]
    return "\n".join(L) + "\n"


def per_class_tables(test: list[Record], by_class: dict[str, tuple[str, str]], reps: dict[str, dict]
                     ) -> tuple[dict, dict]:
    per_class: dict[str, dict[str, dict]] = defaultdict(dict)
    for fname, (src, cls) in by_class.items():
        d = per_class[src].setdefault(cls, {})
        for s, rep in reps.items():
            m = rep.get("by_file", {}).get(fname, {}).get("model", {})
            d["n"] = m.get("n", d.get("n"))
            d[s] = m.get("accuracy")
    n_q = Counter(r.source for r in test)
    per_q = {}
    for src, classes in per_class.items():
        per_q[src] = {"n": n_q[src]}
        for s, rep in reps.items():
            accs = [d[s] for d in classes.values() if d.get(s) is not None]
            per_q[src][s] = {"accuracy": rep.get("by_source", {}).get(src, {}).get("model", {}).get("accuracy"),
                             "balanced_accuracy": sum(accs) / len(accs) if accs else None}
    return dict(per_class), per_q


# ---------------------------------------------------------------- driver

def main(argv: Optional[list[str]] = None) -> dict:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    g = ap.add_argument_group("input")
    g.add_argument("--data", help="labeled .csv / .tsv / .jsonl / .json file")
    g.add_argument("--spec", help="task spec .json / .yaml (see above)")
    g.add_argument("--records", help="already-converted Record JSONL (instead of --data/--spec)")
    g.add_argument("--text-col", help="text column (overrides the spec's text_column; default 'text')")
    g.add_argument("--id-col", help="optional id column")
    g.add_argument("--label-col", action="append", default=[], metavar="QUESTION=COLUMN")
    g.add_argument("--skip-bad-labels", action="store_true", help="drop label values that match no option")
    g = ap.add_argument_group("output / split")
    g.add_argument("--base", default=DEFAULT_BASE, help="checkpoint to adapt (never modified)")
    g.add_argument("--out", required=True, help="output checkpoint dir")
    g.add_argument("--split", default=None, help="train,val,test fractions (default 0.6,0.2,0.2; "
                                                 "calibrate-only 0,0.5,0.5)")
    g.add_argument("--seed", type=int, default=13)
    g.add_argument("--force", action="store_true", help="allow writing into an existing checkpoint dir")
    g.add_argument("--dry-run", action="store_true", help="write the splits and print the plan, run nothing")
    g = ap.add_argument_group("fine-tuning")
    g.add_argument("--calibrate-only", action="store_true", help="skip fine-tuning (for < ~100 examples)")
    g.add_argument("--epochs", type=int, default=3, help="passes over your train split")
    g.add_argument("--replay", default=None, help=f"general Record JSONL to mix in (default {DEFAULT_REPLAY[0]})")
    g.add_argument("--replay-ratio", type=float, default=5.0, help="replay questions per user train question")
    g.add_argument("--lr-backbone", type=float, default=2e-5)
    g.add_argument("--lr-head", type=float, default=1e-4)
    g.add_argument("--warmup", type=float, default=0.1)
    g.add_argument("--grad-accum", type=int, default=4)
    g.add_argument("--max-tokens", type=int, default=2048)
    g.add_argument("--max-batch", type=int, default=64)
    g.add_argument("--max-steps", type=int, default=0, help="0 = about one pass over the mixed train file")
    g.add_argument("--eval-every", type=int, default=0, help="0 = max(5, steps/8)")
    g.add_argument("--max-len", type=int, default=None, help="passed to train/calibrate/evaluate if they accept it")
    g.add_argument("--train-args", default="", help="extra args for jevlike.train, e.g. '--grad-checkpointing'")
    g.add_argument("--keep-last", action="store_true", help="keep OUT/last (resumable optimizer state)")
    g = ap.add_argument_group("calibration / evaluation")
    g.add_argument("--cal-min-count", type=int, default=20, help="val questions needed to fit a temperature cell")
    g.add_argument("--general", default=None, help=f"general held-out JSONL (default {DEFAULT_GENERAL[0]})")
    g.add_argument("--general-eval", type=int, default=500,
                   help="general questions for the forgetting check (and as many again to calibrate "
                        "the cells your data does not cover); 0 = skip")
    g.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto",
                   help="cpu hides the GPU from the subprocesses")
    a = ap.parse_args(argv)
    t_start = time.time()

    base, out = Path(a.base).resolve(), Path(a.out).resolve()
    if not (base / "config.json").exists():
        ap.error(f"--base {base} is not a checkpoint dir (no config.json)")
    if out == base:
        ap.error("--out must differ from --base (the base checkpoint is never modified)")
    if (out / "model.safetensors").exists() and not a.force:
        ap.error(f"{out} already holds a checkpoint; pass --force to overwrite it")
    if (a.records and (a.data or a.spec)) or (not a.records and not (a.data and a.spec)):
        ap.error("give either --data and --spec, or --records")
    work = out / "adapt"
    ddir, edir, ldir = work / "data", work / "eval", work / "logs"
    for d in (ddir, edir, ldir):
        d.mkdir(parents=True, exist_ok=True)
    notes: list[str] = []
    base_cfg = json.loads((base / "config.json").read_text())

    # 1. records + split ------------------------------------------------------------------------
    if a.records:
        recs = read_jsonl(a.records)
        for r in recs:
            r.validate()
        conv = {"rows": None, "skipped": {}, "label_columns": None}
        source_desc, task_desc = f"`{a.records}` (Record JSONL)", {"records": os.path.abspath(a.records)}
    else:
        label_cols = {}
        for kv in a.label_col:
            if "=" not in kv:
                ap.error(f"--label-col wants QUESTION=COLUMN (got {kv!r})")
            k, v = kv.split("=", 1)
            label_cols[k] = v
        spec_d = load_spec_file(a.spec)
        spec = parse_spec(spec_d, a.text_col, a.id_col, label_cols)
        recs, conv = rows_to_records(load_table(a.data), spec, a.skip_bad_labels)
        source_desc = f"`{a.data}` ({conv['rows']} rows)"
        task_desc = {"spec": os.path.abspath(a.spec), "data": os.path.abspath(a.data), "spec_content": spec_d,
                     "text_column": spec.text_column, "label_columns": conv["label_columns"]}
        shutil.copy(a.spec, work / ("task_spec" + Path(a.spec).suffix))
    write_jsonl(str(ddir / "all.jsonl"), recs)
    fractions = parse_split(a.split or ("0,0.5,0.5" if a.calibrate_only else "0.6,0.2,0.2"))
    train, val, test = split_records(recs, fractions, a.seed)
    if not val or not test:
        raise SystemExit(f"split gave val={len(val)} test={len(test)} questions; need more data or another --split")
    if not a.calibrate_only and not train:
        raise SystemExit("the train split is empty; use --calibrate-only or give train a share in --split")
    for name, rs in (("train", train), ("val", val), ("test", test)):
        write_jsonl(str(ddir / f"{name}.jsonl"), rs)
    by_class = write_by_class(test, ddir / "test_by_class")
    n_texts = len({_norm(r.state) for r in recs})
    print(f"{len(recs)} questions over {n_texts} texts -> train {len(train)} / val {len(val)} / test {len(test)}")
    if not a.calibrate_only and n_texts < 100:
        notes.append(f"only {n_texts} labeled texts: fine-tuning is likely noisy; consider --calibrate-only")
        print("WARNING: " + notes[-1])
    if len(val) < a.cal_min_count:
        notes.append(f"val has {len(val)} questions < --cal-min-count {a.cal_min_count}: temperatures keep the fallback")
        print("WARNING: " + notes[-1])

    env = os.environ.copy()
    if a.device == "cpu":
        env["CUDA_VISIBLE_DEVICES"] = ""
    py = sys.executable
    user_states = {_norm(r.state) for r in recs}
    general_file = first_existing([a.general] if a.general else DEFAULT_GENERAL)
    gen_eval = gen_cal = None
    if a.general_eval > 0 and not a.calibrate_only:
        if general_file:
            gs = sample_jsonl(general_file, 2 * a.general_eval, a.seed + 1, user_states)
            gen_eval, gen_cal = ddir / "general_eval.jsonl", ddir / "general_cal.jsonl"
            write_jsonl(str(gen_eval), gs[:a.general_eval])
            write_jsonl(str(gen_cal), gs[a.general_eval:])
            if len(gs) <= a.general_eval:
                gen_cal = None
        else:
            notes.append("no general held-out file found: forgetting check skipped")

    def maxlen_args(module: str) -> list[str]:
        if a.max_len and "--max-len" in cli_flags(module):
            return ["--max-len", str(a.max_len)]
        return []

    def calibrate(ckpt: Path, val_file: Path, fallback: Optional[dict], tag: str) -> dict:
        run_cmd([py, "-m", "jevlike.calibrate", "--ckpt", str(ckpt), "--val", str(val_file),
                 "--min-count", str(a.cal_min_count)] + maxlen_args("jevlike.calibrate"), ldir / f"{tag}.log", env)
        cal = json.loads((ckpt / "calibration.json").read_text())
        return cal if fallback is None else merge_calibration(cal, fallback, a.cal_min_count)

    def evaluate(ckpt: Path, files: list[Path], name: str) -> dict:
        prior = ddir / ("train.jsonl" if train else "val.jsonl")
        run_cmd([py, "-m", "jevlike.evaluate", "--checkpoint", str(ckpt), "--data", *map(str, files),
                 "--out", str(edir), "--name", name, "--latency-calls", "0", "--prior-data", str(prior)]
                + maxlen_args("jevlike.evaluate"), ldir / f"eval_{name}.log", env)
        return json.loads((edir / f"{name}.json").read_text())

    test_files = [ddir / "test_by_class" / f for f in by_class]
    base_cal_path = base / "calibration.json"
    base_cal = json.loads(base_cal_path.read_text()) if base_cal_path.exists() else None

    # 2. plan fine-tuning (written before anything runs, so --dry-run shows it) -----------------
    train_cmd, training = None, None
    if not a.calibrate_only:
        replay_file = first_existing([a.replay] if a.replay else DEFAULT_REPLAY)
        reps = max(1, a.epochs)
        reps_n = int(round(a.replay_ratio * len(train) * reps)) if replay_file else 0
        if a.replay_ratio > 0 and not replay_file:
            notes.append("no replay file found: fine-tuning on your data only (risk of forgetting other questions)")
        excl = {_norm(r.state) for r in val + test}
        replay = sample_jsonl(replay_file, reps_n, a.seed + 2, excl) if reps_n else []
        mix = train * reps + replay
        random.Random(a.seed).shuffle(mix)
        write_jsonl(str(ddir / "train_mix.jsonl"), mix)
        max_len = a.max_len or int(base_cfg.get("max_len", 512))
        steps = a.max_steps or plan_steps(mix, a.max_tokens, a.max_batch, a.grad_accum, max_len)
        eval_every = a.eval_every or max(5, round(steps / 8))
        train_cmd = [py, "-m", "jevlike.train", "--init-from", str(base), "--train", str(ddir / "train_mix.jsonl"),
                     "--val", str(ddir / "val.jsonl"), "--out", str(out), "--epochs", "1", "--max-steps", str(steps),
                     "--eval-every", str(eval_every), "--log-every", str(max(1, min(20, steps // 10))),
                     "--lr-backbone", str(a.lr_backbone), "--lr-head", str(a.lr_head), "--warmup", str(a.warmup),
                     "--grad-accum", str(a.grad_accum), "--max-tokens", str(a.max_tokens),
                     "--max-batch", str(a.max_batch), "--select", "nll", "--seed", str(a.seed)]
        tflags = cli_flags("jevlike.train")
        if "--max-len" in tflags:  # keep the base's context length unless overridden
            train_cmd += ["--max-len", str(max_len)]
        if "--head-max-len" in tflags and "head_max_len" in base_cfg:
            train_cmd += ["--head-max-len", str(base_cfg["head_max_len"])]
        if "large" in str(base_cfg.get("backbone", "")).lower() and "--optim" not in a.train_args:
            # ModernBERT-large only fits a 6 GB GPU with 8-bit optimizer states + checkpointing + frozen bottom layers
            train_cmd += ["--optim", "adamw8bit", "--grad-checkpointing", "--freeze-layers", "10"]
        train_cmd += shlex.split(a.train_args)
        training = {"n_user_train": len(train), "epochs": reps, "n_replay": len(replay),
                    "replay_file": replay_file, "n_mix": len(mix), "max_steps": steps, "eval_every": eval_every,
                    "lr_backbone": a.lr_backbone, "lr_head": a.lr_head, "grad_accum": a.grad_accum,
                    "command": shlex.join(train_cmd)}
        print(f"fine-tune plan: {len(train)} x {reps} + {len(replay)} replay = {len(mix)} questions, "
              f"{steps} steps, eval every {eval_every}")
    if a.dry_run:
        print("dry run: splits written to", ddir)
        if train_cmd:
            print("$ " + shlex.join(train_cmd))
        return {"train": len(train), "val": len(val), "test": len(test), "training": training}

    # 3. base: zero-shot + recalibrated -------------------------------------------------------
    reps_eval: dict[str, dict] = {"base": evaluate(base, test_files, "base_test")}
    general: Optional[dict] = None
    if gen_eval:
        general = {"file": general_file, "n": a.general_eval,
                   "base": metrics_of(evaluate(base, [gen_eval], "base_general"))}
    if a.calibrate_only:
        out.mkdir(parents=True, exist_ok=True)
        for f in base.iterdir():  # a real copy (not links): later training into OUT must not touch the base
            if f.is_file() and f.name not in ("log.jsonl", "train_stdout.txt"):
                shutil.copy2(f, out / ("calibration_base.json" if f.name == "calibration.json" else f.name))
        cfg = json.loads((out / "config.json").read_text())
        if "name" in cfg:
            cfg["name"] = out.name
            (out / "config.json").write_text(json.dumps(cfg, indent=1))
        recal_dir = out
    else:
        recal_dir = work / "base_recal"  # symlinks to the base files + its own calibration.json
        shutil.rmtree(recal_dir, ignore_errors=True)
        recal_dir.mkdir()
        for f in base.iterdir():
            if f.is_file() and not f.name.startswith("calibration"):
                (recal_dir / f.name).symlink_to(f)
    cal = calibrate(recal_dir, ddir / "val.jsonl", base_cal, "calibrate_base_recal")
    (recal_dir / "calibration.json").write_text(json.dumps(cal, indent=1))
    reps_eval["base_recal"] = evaluate(recal_dir, test_files, "base_recal_test")
    final_cal, cal_fallback = cal, "the base checkpoint's calibration.json"

    # 4. fine-tune, calibrate, evaluate --------------------------------------------------------
    if not a.calibrate_only:
        (work / "calibration_base_recal.json").write_text(json.dumps(cal, indent=1))
        shutil.rmtree(recal_dir, ignore_errors=True)
        t0 = time.time()
        run_cmd(train_cmd, ldir / "train.log", env)
        training["seconds"] = round(time.time() - t0, 1)
        if not a.keep_last:
            shutil.rmtree(out / "last", ignore_errors=True)
        fallback = None
        if gen_cal:
            fallback = calibrate(out, gen_cal, None, "calibrate_general")
            (out / "calibration_general.json").write_text(json.dumps(fallback, indent=1))
            cal_fallback = f"fitted on {a.general_eval} general questions (calibration_general.json)"
        else:
            fallback = base_cal
        final_cal = calibrate(out, ddir / "val.jsonl", fallback, "calibrate_adapted")
        (out / "calibration.json").write_text(json.dumps(final_cal, indent=1))
        reps_eval["adapted"] = evaluate(out, test_files, "adapted_test")
        if gen_eval:
            general["adapted"] = metrics_of(evaluate(out, [gen_eval], "adapted_general"))

    # 5. report ---------------------------------------------------------------------------------
    systems = {k: metrics_of(v) for k, v in reps_eval.items()}
    per_class, per_q = per_class_tables(test, by_class, {k: v for k, v in reps_eval.items() if k != "base_recal"})
    if "adapted" in systems:
        b, ad = systems["base"], systems["adapted"]
        if ad["accuracy"] is not None and b["accuracy"] is not None and ad["accuracy"] < b["accuracy"]:
            notes.append("the adapted model is LESS accurate than the base on test: check the labels/spec, "
                         "try more data, a lower LR, or --calibrate-only")
        if general and general.get("adapted") and general["base"]["accuracy"] and \
                general["adapted"]["accuracy"] < general["base"]["accuracy"] - 0.03:
            notes.append("general-data accuracy dropped by > 3 points: raise --replay-ratio or lower the LR")
    if len(test) < 200:
        notes.append(f"the test split has {len(test)} questions, so accuracy differences under ~"
                     f"{196 * math.sqrt(0.25 / len(test)):.0f} points are within noise")
    report = {
        "mode": "calibrate-only" if a.calibrate_only else "fine-tune", "base": str(base), "out": str(out),
        "created": _dt.datetime.now().isoformat(timespec="seconds"), "elapsed_s": round(time.time() - t_start, 1),
        "task": task_desc, "args": vars(a),
        "data": {"source": source_desc, "n_records": len(recs), "n_texts": n_texts, "skipped": conv["skipped"],
                 "fractions": fractions, "split": {"train": len(train), "val": len(val), "test": len(test)},
                 "class_counts": {"train": class_counts(train), "val": class_counts(val), "test": class_counts(test)}},
        "training": training, "systems": systems, "per_question": per_q, "per_class": per_class, "general": general,
        "calibration": {"min_count": a.cal_min_count, "fallback": cal_fallback,
                        "cells": final_cal.get("adapt_merge", {}).get("cells", {}),
                        "temperatures": final_cal.get("temperatures")},
        "notes": notes,
        "files": {"checkpoint": str(out), "splits": str(ddir), "evaluation reports": str(edir), "logs": str(ldir)},
    }
    (out / "adaptation_report.json").write_text(json.dumps(report, indent=1, ensure_ascii=False, default=str))
    md = render_report(report)
    (out / "adaptation_report.md").write_text(md)
    print("\n" + md.split("## Per class")[0].rstrip())
    print(f"\nwrote {out / 'adaptation_report.md'} and .json")
    return report


if __name__ == "__main__":
    main()