Blame for scripts/train_frequency_head.py

"""
Train the Stage 3B frequency classifier on cached freq-v1 features.

Mirrors `scripts/train_head.py` in idiom (torch, seeded, early stop, JSON
report) but operates on the 16-dim frequency features from
`scripts/precompute_frequency_features.py` instead of CLIP embeddings.

Model
-----
A standardised linear logistic-regression head (the "classical first model"
required by S3B-07): feature standardisation (mean/std from train) baked into
the saved checkpoint, then `Linear(16 -> 2)`. Interpretable, tiny, no new
deps (torch only). A small MLP is available via --hidden if the linear model
underperforms, but linear is the default per the tracker's guidance.

Checkpoint format (torch.save dict)
-----------------------------------
  - feature_schema:  "freq-v1"  (must match the features at inference)
  - feature_mean:    float32 (16,)
  - feature_std:     float32 (16,)
  - state_dict:      the classifier head weights
  - hidden:          int (0 = linear logreg)
  - classes:         ["authentic", "ai_generated"]

Usage
-----
    python scripts/train_frequency_head.py \\
        --feat-dir data/freq_features \\
        --out data/checkpoints/freq_v1.pt \\
        --report-out data/reports/freq_v1_train.json
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np

FEATURE_SCHEMA_VERSION = "freq-v1"
_CLASSES = ["authentic", "ai_generated"]


def _load_feat(feat_dir: Path, split: str):
    path = feat_dir / f"{split}_freq.npz"
    with np.load(path, allow_pickle=False) as z:
        schema = str(z["feature_schema"])
        if schema != FEATURE_SCHEMA_VERSION:
            raise ValueError(
                f"{path}: feature_schema {schema!r} != {FEATURE_SCHEMA_VERSION!r}"
            )
        return {
            "features": z["features"].astype(np.float32),
            "labels": z["labels"].astype(np.int64),
            "generators": z["generators"],
            "sources": z["sources"],
        }


def _rates(pred: np.ndarray, y: np.ndarray) -> dict:
    tp = int(((pred == 1) & (y == 1)).sum())
    tn = int(((pred == 0) & (y == 0)).sum())
    fp = int(((pred == 1) & (y == 0)).sum())
    fn = int(((pred == 0) & (y == 1)).sum())
    return {
        "n": int(len(y)),
        "accuracy": float((pred == y).mean()) if len(y) else 0.0,
        "confusion": {"tp": tp, "tn": tn, "fp": fp, "fn": fn},
        "false_positive_rate": float(fp / (fp + tn)) if (fp + tn) else 0.0,
        "false_negative_rate": float(fn / (fn + tp)) if (fn + tp) else 0.0,
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--feat-dir", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--report-out", type=Path, default=None)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--lr", type=float, default=1e-2)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--patience", type=int, default=15)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--hidden",
        type=int,
        default=0,
        help="Hidden units (0 = linear logistic regression).",
    )
    parser.add_argument(
        "--eval-split",
        action="append",
        default=[],
        help="Extra split(s) to evaluate, e.g. --eval-split heldout.",
    )
    args = parser.parse_args()

    import torch
    import torch.nn as nn

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    train = _load_feat(args.feat_dir, "train")
    val = _load_feat(args.feat_dir, "val")
    test = _load_feat(args.feat_dir, "test")
    extra = {name: _load_feat(args.feat_dir, name) for name in args.eval_split}

    # Standardisation from train only (baked into the checkpoint).
    mean = train["features"].mean(axis=0)
    std = train["features"].std(axis=0) + 1e-8

    def prep(split):
        x = (split["features"] - mean) / std
        return torch.from_numpy(x.astype(np.float32)), torch.from_numpy(
            split["labels"]
        )

    x_tr, y_tr = prep(train)
    x_val, y_val = prep(val)

    n_feat = x_tr.shape[1]
    if args.hidden > 0:
        head = nn.Sequential(
            nn.Linear(n_feat, args.hidden),
            nn.ReLU(),
            nn.Linear(args.hidden, 2),
        )
    else:
        head = nn.Linear(n_feat, 2)

    # Class weights to counter the authentic/ai imbalance in the splits.
    counts = np.bincount(train["labels"], minlength=2).astype(np.float64)
    w = counts.sum() / (2.0 * np.maximum(counts, 1.0))
    class_weight = torch.tensor(w, dtype=torch.float32)
    loss_fn = nn.CrossEntropyLoss(weight=class_weight)
    opt = torch.optim.Adam(
        head.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )

    print(
        f"train={tuple(x_tr.shape)} val={tuple(x_val.shape)} "
        f"class_counts={counts.tolist()} hidden={args.hidden}"
    )

    best_val = float("inf")
    best_state = None
    no_improve = 0
    for epoch in range(1, args.epochs + 1):
        head.train()
        perm = torch.randperm(len(x_tr))
        losses = []
        for i in range(0, len(x_tr), args.batch_size):
            idx = perm[i : i + args.batch_size]
            opt.zero_grad()
            loss = loss_fn(head(x_tr[idx]), y_tr[idx])
            loss.backward()
            opt.step()
            losses.append(loss.item())
        head.eval()
        with torch.no_grad():
            vloss = loss_fn(head(x_val), y_val).item()
            vacc = float((head(x_val).argmax(-1) == y_val).float().mean())
        if epoch % 20 == 0 or epoch == 1:
            print(
                f"  epoch {epoch:3d} train_loss={np.mean(losses):.4f} "
                f"val_loss={vloss:.4f} val_acc={vacc:.4f}"
            )
        if vloss < best_val - 1e-4:
            best_val = vloss
            best_state = {k: v.clone() for k, v in head.state_dict().items()}
            no_improve = 0
        else:
            no_improve += 1
            if no_improve >= args.patience:
                print(f"  early stop at epoch {epoch}")
                break

    if best_state is not None:
        head.load_state_dict(best_state)
    head.eval()

    def evaluate(split) -> dict:
        x = (split["features"] - mean) / std
        with torch.no_grad():
            logits = head(torch.from_numpy(x.astype(np.float32)))
            pred = logits.argmax(-1).numpy()
        y = split["labels"]
        rep = _rates(pred, y)
        # Per-generator breakdown over AI rows.
        gens = np.asarray([str(g) for g in split["generators"]])
        by_gen = {}
        for g in sorted(set(gens) - {""}):
            m = gens == g
            by_gen[g] = _rates(pred[m], y[m])
        if by_gen:
            rep["by_generator"] = by_gen
        return rep

    report = {
        "feature_schema": FEATURE_SCHEMA_VERSION,
        "hidden": args.hidden,
        "splits": {
            "val": evaluate(val),
            "test": evaluate(test),
            **{name: evaluate(s) for name, s in extra.items()},
        },
    }
    for split_name in ("test", *args.eval_split):
        r = report["splits"][split_name]
        print(
            f"  [{split_name}] acc={r['accuracy']:.4f} "
            f"fpr={r['false_positive_rate']:.4f} fnr={r['false_negative_rate']:.4f}"
        )

    args.out.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "feature_schema": FEATURE_SCHEMA_VERSION,
            "feature_mean": mean.astype(np.float32),
            "feature_std": std.astype(np.float32),
            "state_dict": head.state_dict(),
            "hidden": args.hidden,
            "classes": _CLASSES,
        },
        args.out,
    )
    print(f"\nsaved frequency head to {args.out}")

    if args.report_out is not None:
        args.report_out.parent.mkdir(parents=True, exist_ok=True)
        with args.report_out.open("w", encoding="utf-8") as fh:
            json.dump(report, fh, indent=2, sort_keys=True)
        print(f"saved report to {args.report_out}")


if __name__ == "__main__":
    main()