scripts/03_train.py

"""Fine-tune an encoder and evaluate it on the held-out test set.

    python scripts/03_train.py --model modernbert --scheme document
    python scripts/03_train.py --model modernbert --scheme random
"""

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from torch.utils.data import DataLoader  # noqa: E402

from cti_attack import config, data, evaluate, modeling  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=config.DEFAULT_MODEL, choices=list(config.BASE_MODELS))
    ap.add_argument("--scheme", default="document", choices=["document", "random"])
    ap.add_argument("--epochs", type=int, default=config.EPOCHS)
    args = ap.parse_args()

    labels = data.load_labels()
    tr = data.load_split(args.scheme, "train")
    dv = data.load_split(args.scheme, "dev")
    te = data.load_split(args.scheme, "test")

    run = f"{args.model}__{args.scheme}"
    out_dir = config.MODELS_DIR / run
    print(f"{'=' * 62}\n  {run}\n{'=' * 62}")
    print(f"  train={len(tr)}  dev={len(dv)}  test={len(te)}  labels={len(labels)}")

    best = modeling.train(args.model, args.scheme, tr, dv, labels, out_dir, epochs=args.epochs)

    # ---- evaluate the best checkpoint on dev (tuning) then test (reporting) --
    model, tok = modeling.load_for_inference(out_dir)
    dev_ = modeling.device()
    model.to(dev_)

    amp = args.model not in config.FP32_ONLY_MODELS
    ds_dv = modeling.SentenceDataset(dv, labels, tok, config.MAX_LENGTH)
    ds_te = modeling.SentenceDataset(te, labels, tok, config.MAX_LENGTH)
    s_dv = modeling.predict_scores(model, DataLoader(ds_dv, batch_size=32), dev_, amp=amp)
    s_te = modeling.predict_scores(model, DataLoader(ds_te, batch_size=32), dev_, amp=amp)

    Ydv = ds_dv.Y.astype("int8")
    Yte = ds_te.Y.astype("int8")

    gt, _ = evaluate.tune_global_threshold(Ydv, s_dv)
    pct = evaluate.tune_per_class_thresholds(Ydv, s_dv)

    rep_g = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, gt), labels)
    rep_p = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, pct), labels)

    print(f"\n  TEST global t={gt}   macro-F1={rep_g.macro_f1:.4f}  micro-F1={rep_g.micro_f1:.4f}")
    print(f"  TEST per-class       macro-F1={rep_p.macro_f1:.4f}  micro-F1={rep_p.micro_f1:.4f}")

    payload = {
        "run": run,
        "base_model": config.BASE_MODELS[args.model],
        "split_scheme": args.scheme,
        "best_epoch": best["epoch"],
        "dev_macro_f1": round(best["macro_f1"], 4),
        "global_threshold": {"threshold": gt, **rep_g.as_dict()},
        "per_class_threshold": {
            "thresholds": {l: float(t) for l, t in zip(labels, pct)},
            **rep_p.as_dict(),
        },
    }
    evaluate.save_report(args.model, args.scheme, payload)
    (out_dir / "thresholds.json").write_text(
        json.dumps({"global": gt,
                    "per_class": {l: float(t) for l, t in zip(labels, pct)}}, indent=2),
        encoding="utf-8")


if __name__ == "__main__":
    main()