Spaces:

Veridicate
/

scanner

Running

scanner / scripts /train_frequency_head.py

github-actions[bot]

Deploy from GitHub 1af737237f5d10dfa9cf12849d98f47c3b30ce37 (filtered: code + Dockerfile + README + NOTICES only)

a2f55b2 about 2 months ago

8.31 kB

"""
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()

	"""
	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()