code/minijev/training/supervised.py

"""Supervised / distillation training loop.

Resumable (last.pt holds model, optimizer, scheduler, scaler, data position and
RNG), stops cleanly on SIGINT/SIGTERM or a time budget, and ends with
temperature calibration on the validation split and a test-set report.
"""
from __future__ import annotations

import math
import os
import random
import signal
import sys
import time
from pathlib import Path

import numpy as np
import torch

from minijev.calibration.metrics import (IsotonicConfidence, calibration_report, fit_temperature,
                                         reliability_svg, softmax)
from minijev.datasets.batching import BackboneTable, BatchEncoder, EpochBatches, model_inputs, to_torch
from minijev.datasets.decision_data import DecisionData
from minijev.device import RuntimeProfile, autocast, resolve_runtime
from minijev.model import build_model, human_params
from minijev.runs import RunLogger, environment_info, new_run_dir
from minijev.tokenizers import make_tokenizer
from minijev.tasks import get_task

from .losses import decision_loss


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def task_specs(cfg: dict) -> list[dict]:
    """Single task (`task:` + `data:`) or multi-task (`tasks: [{key, task, data, weight}]`)."""
    if cfg.get("tasks"):
        sampling = cfg.get("sampling") or {}
        out = []
        for t in cfg["tasks"]:
            t = dict(t)
            t.setdefault("key", t["task"]["name"])
            t["weight"] = float(sampling.get(t["key"], t.get("weight", 1.0)))
            out.append(t)
        return out
    return [{"key": cfg["task"]["name"], "task": cfg["task"], "data": cfg["data"], "weight": 1.0}]


def make_task_for(cfg: dict, task_cfg: dict):
    """Task plugin with the model's text tokenizer attached (byte or pretrained)."""
    task = get_task(task_cfg["name"], task_cfg)
    task.tokenizer = make_tokenizer(cfg["model"].get("text_encoder"))
    return task


def load_splits(cfg: dict) -> dict[str, DecisionData]:
    dcfg = cfg["data"]
    d = Path(dcfg["dir"])
    if not (d / "train.npz").exists():
        gen = dcfg.get("generate")
        if not gen:
            raise FileNotFoundError(f"{d}/train.npz missing and no data.generate section")
        import importlib
        mod = importlib.import_module(dcfg.get("generator") or f"environments.{cfg['task']['name']}.dataset")
        gen = dict(gen)
        n = gen.pop("samples")
        print(f"[data] generating {n:,} samples into {d} ...")
        mod.generate_splits(d, n, **gen)
    splits = {k: DecisionData.load(d / f"{k}.npz") for k in ("train", "val", "test")}
    lim = dcfg.get("train_limit")
    if lim and lim < len(splits["train"]):
        splits["train"] = splits["train"].subset(np.arange(lim))
    return splits


def make_model(cfg: dict, task=None):
    """task: one task, a list of tasks (multi-task: one state encoder per modality), or None (from cfg)."""
    tasks = task if isinstance(task, (list, tuple)) else [task] if task is not None else \
        [make_task_for(cfg, t["task"]) for t in task_specs(cfg)]
    encs, classes = {}, {}
    for t in tasks:
        encs.setdefault(t.state_encoder, t.encoder_config())
        if t.encoder_class():
            classes[t.state_encoder] = t.encoder_class()
    return build_model(cfg["model"], encs, classes or None)


def param_groups(model, wd: float, lr: float | None = None, backbone_lr: float | None = None):
    """AdamW groups: no decay on norms/biases/embeddings; a pretrained backbone can use its own (lower) lr."""
    groups = {}
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        nd = p.ndim < 2 or "embed" in n or n.endswith((".pos", ".row", ".col", "cls", "type_emb"))
        bb = backbone_lr is not None and n.startswith("backbone.")
        groups.setdefault((nd, bb), []).append(p)
    out = []
    for (nd, bb), ps in groups.items():
        g = {"params": ps, "weight_decay": 0.0 if nd else wd}
        if bb:
            g["lr"] = backbone_lr
        elif lr is not None:
            g["lr"] = lr
        out.append(g)
    return out


def lr_lambda(warmup: int, total: int, min_ratio: float):
    def f(step):
        if step < warmup:
            return (step + 1) / max(1, warmup)
        prog = min(1.0, (step - warmup) / max(1, total - warmup))
        return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * prog))
    return f


@torch.no_grad()
def predict_logits(model, encoder: BatchEncoder, device, batch_size: int = 1024, max_rows: int | None = None,
                   profile: RuntimeProfile | None = None) -> tuple[np.ndarray, np.ndarray]:
    model.eval()
    n = len(encoder.data) if max_rows is None else min(max_rows, len(encoder.data))
    outs, tg = [], []
    for i in range(0, n, batch_size):
        b = to_torch(encoder(np.arange(i, min(n, i + batch_size))), device)
        with autocast(profile) if profile else torch.autocast("cpu", enabled=False):
            out = model(**model_inputs(b))
        outs.append(out["logits"].float().cpu().numpy())
        tg.append(b["target"].cpu().numpy())
    return np.concatenate(outs), np.concatenate(tg)


def _acceptable(encoder: BatchEncoder, n: int):
    """Rows' acceptable choices (soft label > 0), if the dataset has soft labels."""
    soft = encoder.data.soft
    return None if soft is None else soft[:n] > 0


class _Stop:
    def __init__(self):
        self.flag = False
        self._prev = {}

    def __enter__(self):
        for s in (signal.SIGINT, signal.SIGTERM):
            try:
                self._prev[s] = signal.signal(s, self._handle)
            except ValueError:  # not in main thread
                pass
        return self

    def _handle(self, *_):
        if self.flag:
            raise KeyboardInterrupt
        print("\n[train] stop requested - saving checkpoint (Ctrl-C again to abort)")
        self.flag = True

    def __exit__(self, *a):
        for s, h in self._prev.items():
            signal.signal(s, h)


class MultiTaskBatches:
    """Batches for multi-task training: each batch comes from ONE task (its own state
    encoder / choice set), picked with the configured sampling weights. Deterministic
    in (seed, epoch, i), so resume skips exactly the consumed batches."""

    def __init__(self, encoders: dict, weights: dict, batch_size: int, n_batches: int, seed: int, epoch: int,
                 start_batch: int = 0):
        self.keys = list(encoders)
        w = np.array([weights[k] for k in self.keys], float)
        self.p = w / w.sum()
        self.encoders, self.bs, self.seed, self.epoch = encoders, batch_size, seed, epoch
        self.range = range(start_batch, n_batches)

    def __len__(self):
        return len(self.range)

    def __getitem__(self, i):
        rng = np.random.default_rng((self.seed, self.epoch, self.range[i], 11))
        enc = self.encoders[self.keys[rng.choice(len(self.keys), p=self.p)]]
        n = len(enc.data)
        return enc(rng.choice(n, self.bs, replace=n < self.bs), rng)


def make_encoders(cfg: dict, model, task, split: DecisionData, choice_dropout: float = 0.0, seed: int = 0):
    """BatchEncoder for one split; with a frozen pretrained backbone, questions and
    choice labels are encoded once per vocabulary entry (BackboneTable)."""
    mcfg = cfg["model"]
    kw = dict(max_question_len=mcfg.get("max_question_len", 32), max_choice_len=mcfg.get("max_choice_len", 16))
    tables = {}
    bb = getattr(model, "backbone", None)
    if bb is not None and bb.frozen and cfg["training"].get("cache_text", True):
        tok = task.tokenizer
        tables = {"question_table": BackboneTable(split.questions, model, tok, kw["max_question_len"]),
                  "choice_table": BackboneTable([task.model_label(c) for c in split.choices], model, tok,
                                                kw["max_choice_len"])}
    return BatchEncoder(split, task, choice_dropout=choice_dropout, seed=seed, tokenizer=task.tokenizer,
                        **kw, **tables)


def train(cfg: dict, run_dir: str | Path | None = None, resume: bool = False, max_steps: int | None = None,
          max_minutes: float | None = None, quiet: bool = False, final_eval: bool = True) -> dict:
    tcfg = cfg["training"]
    seed = int(cfg.get("seed", 0))
    seed_everything(seed)
    specs = task_specs(cfg)
    tasks = {s["key"]: make_task_for(cfg, s["task"]) for s in specs}
    splits = {s["key"]: load_splits({**cfg, "task": s["task"], "data": s["data"]}) for s in specs}
    weights = {s["key"]: s["weight"] for s in specs}
    multi = len(specs) > 1
    model = make_model(cfg, list(tasks.values()))
    if tcfg.get("init_checkpoint") and not resume:
        ck = torch.load(tcfg["init_checkpoint"], map_location="cpu", weights_only=False)
        model.load_state_dict(ck["model"])
        model.log_temperature.zero_()  # re-calibrated at the end of this run
        print(f"[train] initialised from {tcfg['init_checkpoint']}")
    n_params = model.num_parameters()
    prof = resolve_runtime(tcfg, n_params)
    torch.set_num_threads(prof.threads)
    device = torch.device(prof.device)
    model.to(device)

    run_dir = Path(run_dir) if run_dir else new_run_dir(cfg.get("name", "run"))
    ckpt_dir = run_dir / "checkpoints"
    log = RunLogger(run_dir, cfg, resume=resume)
    encs = {}
    for k, sp in splits.items():
        encs[k] = {"train": make_encoders(cfg, model, tasks[k], sp["train"], tcfg.get("choice_dropout", 0.0), seed),
                   "val": make_encoders(cfg, model, tasks[k], sp["val"]),
                   "test": make_encoders(cfg, model, tasks[k], sp["test"])}
    train_enc = {k: e["train"] for k, e in encs.items()}

    eff_bs = int(tcfg.get("batch_size", 256))
    epochs = float(tcfg.get("epochs", 1))
    n_train = sum(len(sp["train"]) for sp in splits.values())
    batches_per_epoch = max(1, n_train // prof.micro_batch)
    steps_per_epoch = max(1, batches_per_epoch // prof.grad_accum)
    total_steps = max(1, int(steps_per_epoch * epochs))
    if max_steps:
        total_steps = min(total_steps, max_steps)
    lr = tcfg.get("lr", 1e-3)
    opt = torch.optim.AdamW(param_groups(model, tcfg.get("weight_decay", 0.05), lr, tcfg.get("backbone_lr")), lr=lr,
                            betas=(0.9, 0.98), fused=prof.device == "cuda")
    warmup = int(tcfg.get("warmup_steps", min(500, total_steps // 20)))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(warmup, total_steps, tcfg.get("min_lr_ratio", 0.05)))
    scaler = torch.amp.GradScaler(enabled=prof.amp_dtype == "fp16")
    fwd = torch.compile(model) if prof.compile else model

    step, epoch, batch_pos, best, elapsed_prev = 0, 0, 0, -1.0, 0.0
    samples_seen = 0
    session_samples = 0
    last = ckpt_dir / "last.pt"
    if resume and last.exists():
        ck = torch.load(last, map_location=device, weights_only=False)
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["optimizer"])
        sched.load_state_dict(ck["scheduler"])
        scaler.load_state_dict(ck["scaler"])
        step, epoch, batch_pos, best = ck["step"], ck["epoch"], ck["batch_pos"], ck["best"]
        elapsed_prev, samples_seen = ck.get("elapsed", 0.0), ck.get("samples_seen", 0)
        torch.set_rng_state(ck["torch_rng"])
        print(f"[train] resumed from {last} at step {step} (epoch {epoch}, batch {batch_pos})")

    breakdown = model.param_breakdown()
    info = {"task": "+".join(tasks) if multi else next(iter(tasks.values())).name, "tasks": list(tasks),
            "params": n_params, "params_human": human_params(n_params), "params_breakdown": breakdown,
            "train_rows": n_train, "val_rows": sum(len(sp["val"]) for sp in splits.values()),
            "test_rows": sum(len(sp["test"]) for sp in splits.values()),
            "effective_batch": eff_bs, "total_steps": total_steps, "runtime": prof.to_dict(),
            "env": environment_info()}
    log.write_json("env.json", info)
    if not quiet:
        print(f"[train] {cfg.get('name')} | {human_params(n_params)} params (core+heads "
              f"{human_params(breakdown['core_total'])}, text {human_params(breakdown['text'])}, trainable "
              f"{human_params(breakdown['trainable'])}) | {prof.device} ({prof.device_name}) "
              f"amp={prof.amp_dtype} micro={prof.micro_batch}x{prof.grad_accum} | "
              f"{n_train:,} train rows ({', '.join(tasks)}) | {total_steps:,} steps")

    def save(path, extra=None):
        tmp = Path(str(path) + ".tmp")
        torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "scheduler": sched.state_dict(),
                    "scaler": scaler.state_dict(), "step": step, "epoch": epoch, "batch_pos": batch_pos,
                    "best": best, "elapsed": elapsed_prev + (time.time() - t_start), "samples_seen": samples_seen,
                    "torch_rng": torch.get_rng_state(), "config": cfg, **(extra or {})}, tmp)
        os.replace(tmp, path)

    def validate():
        reps = {}
        for k, e in encs.items():
            lg, tg = predict_logits(model, e["val"], device, max_rows=eval_rows, profile=prof)
            reps[k] = calibration_report(lg, tg, acceptable=_acceptable(e["val"], len(tg)))
        w = np.array([weights[k] for k in reps])
        acc = float(np.dot(w / w.sum(), [r["accuracy"] for r in reps.values()]))
        return acc, reps

    loss_cfg = cfg.get("loss", {})
    eval_every = int(tcfg.get("eval_every", 1000))
    ckpt_every = int(tcfg.get("checkpoint_every", 1000))
    log_every = int(tcfg.get("log_every", 50))
    eval_rows = tcfg.get("eval_max_rows", 20000)
    grad_clip = tcfg.get("grad_clip", 1.0)
    t_start = time.time()
    t_train = 0.0
    run_loss, run_n, window_samples, window_t = 0.0, 0, 0, time.time()
    stopped_early = False

    with _Stop() as stop:
        while step < total_steps and not stop.flag:
            if multi:
                ds = MultiTaskBatches(train_enc, weights, prof.micro_batch, batches_per_epoch, seed, epoch, batch_pos)
            else:
                ds = EpochBatches(next(iter(train_enc.values())), prof.micro_batch, shuffle=True, seed=seed,
                                  epoch=epoch, start_batch=batch_pos, drop_last=True)
            loader = torch.utils.data.DataLoader(ds, batch_size=None, shuffle=False, num_workers=prof.num_workers,
                                                 pin_memory=prof.device == "cuda",
                                                 prefetch_factor=4 if prof.num_workers else None,
                                                 # fork shares the in-memory dataset with workers instead of pickling it
                                                 multiprocessing_context="fork" if prof.num_workers and sys.platform == "linux" else None)
            model.train()
            micro = 0
            for batch in loader:
                t0 = time.time()
                b = to_torch(batch, device, non_blocking=True)
                with autocast(prof):
                    out = fwd(**model_inputs(b))
                loss, terms = decision_loss(out["logits"], out["choice_mask"], b, loss_cfg, out.get("score"))
                scaler.scale(loss / prof.grad_accum).backward()
                micro += 1
                batch_pos += 1
                n_b = len(batch["target"])
                samples_seen += n_b
                session_samples += n_b
                window_samples += n_b
                run_loss += float(loss.detach())
                run_n += 1
                if micro % prof.grad_accum:
                    t_train += time.time() - t0
                    continue
                if grad_clip:
                    scaler.unscale_(opt)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(opt)
                scaler.update()
                opt.zero_grad(set_to_none=True)
                sched.step()
                step += 1
                t_train += time.time() - t0
                if step % log_every == 0:
                    now = time.time()
                    sps = window_samples / max(1e-9, now - window_t)
                    log.log(step, train_loss=run_loss / run_n, lr=sched.get_last_lr()[0],
                            samples_per_sec=round(sps, 1), epoch=epoch, **{f"loss_{k}": v for k, v in terms.items()})
                    if not quiet:
                        print(f"  step {step:6d}/{total_steps} loss {run_loss / run_n:.4f} "
                              f"lr {sched.get_last_lr()[0]:.2e} {sps:,.0f} samples/s")
                    run_loss, run_n, window_samples, window_t = 0.0, 0, 0, now
                if step % eval_every == 0 or step == total_steps:
                    acc, reps = validate()
                    first = next(iter(reps.values()))
                    extra = {"val_acceptable_acc": first["acceptable_accuracy"]} \
                        if not multi and "acceptable_accuracy" in first else {}
                    per_task = {f"val_acc_{k}": r["accuracy"] for k, r in reps.items()} if multi else {}
                    log.log(step, val_acc=acc, val_nll=first["nll"], val_ece=first["ece"],
                            val_brier=first["brier"], **extra, **per_task)
                    if not quiet:
                        detail = " ".join(f"{k} {r['accuracy']:.4f}" for k, r in reps.items()) if multi else \
                            f"nll {first['nll']:.4f} ece {first['ece']:.4f}"
                        print(f"  [val] step {step} acc {acc:.4f} {detail}"
                              + (f" acceptable {first['acceptable_accuracy']:.4f}" if extra else ""))
                    if acc > best:
                        best = acc
                        torch.save({"model": model.state_dict(), "config": cfg, "step": step,
                                    "val_accuracy": best}, ckpt_dir / "best.pt")
                    model.train()
                if step % ckpt_every == 0:
                    save(last)
                if max_minutes and (time.time() - t_start) / 60 > max_minutes:
                    print(f"[train] time budget of {max_minutes} min reached")
                    stopped_early = True
                    stop.flag = True
                if step >= total_steps or stop.flag:
                    break
            else:
                epoch += 1
                batch_pos = 0
                continue
            if batch_pos >= batches_per_epoch:
                epoch, batch_pos = epoch + 1, 0
        save(last)
        interrupted = stop.flag and not stopped_early and step < total_steps

    train_seconds = elapsed_prev + (time.time() - t_start)
    summary = {**{k: v for k, v in info.items() if k != "env"}, "steps_done": step, "samples_seen": samples_seen,
               "train_seconds": round(train_seconds, 1),
               "throughput_samples_per_sec": round(session_samples / max(t_train, 1e-9), 1),
               "best_val_accuracy": best, "complete": step >= total_steps, "run_dir": str(run_dir)}
    hourly = cfg.get("cost", {}).get("gpu_hourly_cost") or os.environ.get("MINIJEV_GPU_HOURLY_COST")
    if hourly:
        summary["cost"] = {"gpu_hourly_cost": float(hourly),
                           "estimated_cost": round(float(hourly) * train_seconds / 3600, 4)}
    if final_eval and not interrupted:
        summary.update(finalize(cfg, model, run_dir, {k: e["val"] for k, e in encs.items()},
                                {k: e["test"] for k, e in encs.items()}, device, prof))
    log.write_json("summary.json", summary)
    log.close()
    return summary


def _pad_logits(parts: list[np.ndarray]) -> np.ndarray:
    n = max(p.shape[1] for p in parts)
    return np.concatenate([np.pad(p, ((0, 0), (0, n - p.shape[1])), constant_values=-1e9) for p in parts])


def finalize(cfg, model, run_dir: Path, val_encs, test_encs, device, prof) -> dict:
    """Best checkpoint -> temperature fitted on val (global + per task) -> calibrated test report -> final.pt."""
    from minijev.calibration.metrics import ece as ece_fn

    if not isinstance(val_encs, dict):  # single-encoder call (old signature)
        val_encs, test_encs = {"task": val_encs}, {"task": test_encs}
    ckpt_dir = run_dir / "checkpoints"
    best = ckpt_dir / "best.pt"
    if best.exists():
        model.load_state_dict(torch.load(best, map_location=device, weights_only=False)["model"])
    cal = cfg.get("calibration", {})
    per_task, V, VT, TL, TT = {}, [], [], [], []
    for k in val_encs:
        vl, vt = predict_logits(model, val_encs[k], device, profile=prof)
        tl, tt = predict_logits(model, test_encs[k], device, profile=prof)
        Tk = fit_temperature(vl, vt) if cal.get("method", "temperature") == "temperature" else 1.0
        acc_ok = _acceptable(test_encs[k], len(tt))
        per_task[k] = {"temperature": Tk, "test_raw": calibration_report(tl, tt, 1.0, acceptable=acc_ok),
                       "test_calibrated": calibration_report(tl, tt, Tk, acceptable=acc_ok)}
        V.append(vl), VT.append(vt), TL.append(tl), TT.append(tt)
    val_lg, val_tg, test_lg, test_tg = _pad_logits(V), np.concatenate(VT), _pad_logits(TL), np.concatenate(TT)
    T = fit_temperature(val_lg, val_tg) if cal.get("method", "temperature") == "temperature" else 1.0
    model.log_temperature.fill_(math.log(T))
    single = len(val_encs) == 1
    acc_ok = _acceptable(next(iter(test_encs.values())), len(test_tg)) if single else None
    raw = calibration_report(test_lg, test_tg, 1.0, acceptable=acc_ok)
    calibrated = calibration_report(test_lg, test_tg, T, acceptable=acc_ok)
    # isotonic comparison (fit on val confidence, evaluate on test)
    pv = softmax(val_lg, T)
    iso = IsotonicConfidence().fit(pv.max(1), (pv.argmax(1) == val_tg).astype(float))
    pt = softmax(test_lg, T)
    iso_ece = ece_fn(iso(pt.max(1)), (pt.argmax(1) == test_tg).astype(float))
    (run_dir / "reliability_raw.svg").write_text(reliability_svg(raw["reliability"], "Test - raw"))
    (run_dir / "reliability_calibrated.svg").write_text(
        reliability_svg(calibrated["reliability"], f"Test - temperature {T:.3f}"))
    task_T = {k: v["temperature"] for k, v in per_task.items()}
    torch.save({"model": model.state_dict(), "config": cfg, "temperature": T, "task_temperatures": task_T},
               ckpt_dir / "final.pt")
    print(f"[calibration] T={T:.3f} | test acc {calibrated['accuracy']:.4f} | ECE {raw['ece']:.4f} -> "
          f"{calibrated['ece']:.4f} | NLL {raw['nll']:.4f} -> {calibrated['nll']:.4f} | "
          f"Brier {calibrated['brier']:.4f} | isotonic ECE {iso_ece:.4f}"
          + (f" | acceptable acc {calibrated['acceptable_accuracy']:.4f}" if acc_ok is not None else ""))
    if not single:
        for k, v in per_task.items():
            c = v["test_calibrated"]
            print(f"  [{k}] T={v['temperature']:.3f} acc {c['accuracy']:.4f} ECE {v['test_raw']['ece']:.4f} -> "
                  f"{c['ece']:.4f} NLL {c['nll']:.4f}")
    return {"temperature": T, "task_temperatures": task_T, "test_raw": raw, "test_calibrated": calibrated,
            "test_isotonic_ece": iso_ece, "per_task": per_task, "final_checkpoint": str(ckpt_dir / "final.pt")}