""" AgentSight — full training entry point. Key design decisions ──────────────────── * AdamW with cosine-annealing LR schedule (linear warmup for 10% of steps). * Gradient accumulation over grad_accum_steps=4 trajectories (effective batch ~4). * pos_weight computed dynamically from the training split ratio. * WeightedRandomSampler ensures each epoch sees balanced class distribution. * Decision threshold tuned on val set after training (never on test). * Best model selected by val step_localization_accuracy (primary metric). * Test set NEVER evaluated during training — locked behind hash verification. TEST SET INTEGRITY SEAL — DO NOT MODIFY sha256sum data/splits/test.json 9604aae8eb5aec4ae666cfbe3053910f0570a807a4fa5515223dbca1aa66a7d8 test.json is LOCKED until a final deliberate single run for the paper. """ import os import sys import json import argparse script_dir = os.path.dirname(os.path.abspath(__file__)) project_root = os.path.join(script_dir, "..", "..") sys.path.insert(0, project_root) import torch import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR from src.data.preprocessor import StepPreprocessor from src.data.dataset import get_dataloader, AgentTrajectoryDataset from src.models.agentsight import AgentSightModel from src.training.train import train_epoch, compute_pos_weight from src.training.evaluate import evaluate, tune_threshold def main(): parser = argparse.ArgumentParser(description="Train AgentSight Hallucination Detector") parser.add_argument("--epochs", type=int, default=50, help="Maximum training epochs") parser.add_argument("--lr", type=float, default=3e-5, help="Peak learning rate for AdamW") parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps") parser.add_argument("--patience", type=int, default=15, help="Early stopping patience (epochs)") parser.add_argument("--max_len", type=int, default=512, help="Tokenizer max sequence length") parser.add_argument("--max_steps", type=int, default=20, help="Max trajectory steps (centred truncation)") parser.add_argument("--weight_decay", type=float, default=0.01, help="AdamW weight decay") parser.add_argument("--warmup_ratio", type=float, default=0.10, help="Fraction of steps for LR warmup") parser.add_argument("--no_weighted_sampler", action="store_true", help="Disable WeightedRandomSampler") args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") splits_dir = os.path.join(project_root, "data", "splits") # ── Preprocessor & data ─────────────────────────────────────────────────── print("Loading tokenizer …") preprocessor = StepPreprocessor(max_len=args.max_len) use_sampler = not args.no_weighted_sampler print(f"Building dataloaders (WeightedRandomSampler={use_sampler}) …") train_loader = get_dataloader( os.path.join(splits_dir, "train.json"), preprocessor, batch_size=1, shuffle=True, use_weighted_sampler=use_sampler, ) val_loader = get_dataloader( os.path.join(splits_dir, "val.json"), preprocessor, batch_size=1, shuffle=False, ) with open(os.path.join(splits_dir, "val.json")) as f: val_samples = json.load(f) print("Initialising AgentSightModel …") model = AgentSightModel() model.to(device) # ── Optimiser & scheduler ───────────────────────────────────────────────── # Separate LoRA params from the rest to avoid weight-decaying bias/LN terms no_decay = ["bias", "LayerNorm.weight", "LayerNorm.bias"] param_groups = [ { "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": args.weight_decay, }, { "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] optimizer = optim.AdamW(param_groups, lr=args.lr) total_updates = (len(train_loader) // args.grad_accum) * args.epochs warmup_steps = int(total_updates * args.warmup_ratio) # Linear warmup then cosine decay implemented manually via LambdaLR def lr_lambda(step): if step < warmup_steps: return float(step) / max(1, warmup_steps) progress = float(step - warmup_steps) / max(1, total_updates - warmup_steps) import math return 0.5 * (1.0 + math.cos(math.pi * progress)) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) # ── Dynamic pos_weight ──────────────────────────────────────────────────── print("Computing class weights from training data …") pos_weight = compute_pos_weight(train_loader, device) print(f" pos_weight = {pos_weight.item():.2f}") # ── Training loop ───────────────────────────────────────────────────────── weights_path = os.path.join(project_root, "src", "models", "best_agentsight.pth") best_val_step_acc = 0.0 best_threshold = 0.5 epochs_no_improve = 0 print(f"\nStarting training for up to {args.epochs} epochs " f"(patience={args.patience}) …\n") for epoch in range(1, args.epochs + 1): print(f"── Epoch {epoch}/{args.epochs} ──────────────────────────────") avg_loss = train_epoch( model, train_loader, optimizer, scheduler, pos_weight, grad_accum_steps=args.grad_accum, ) print(f" Train loss : {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.2e}") # Tune threshold on val every 5 epochs (cheap) or when we might improve if epoch % 5 == 0 or epoch <= 5: thr, thr_f1 = tune_threshold(model, val_samples, preprocessor) print(f" Val threshold tuning → thr={thr:.2f} macro-F1={thr_f1*100:.1f}%") else: thr = best_threshold metrics = evaluate(model, val_samples, preprocessor, threshold=thr) print( f" Val step-acc : {metrics['step_acc']*100:.1f}% | " f"F1 : {metrics['judgment_f1']*100:.1f}% | " f"Recall : {metrics['judgment_recall']*100:.1f}% | " f"Precision : {metrics['judgment_precision']*100:.1f}%" ) if metrics["step_acc"] > best_val_step_acc: best_val_step_acc = metrics["step_acc"] best_threshold = thr epochs_no_improve = 0 torch.save(model.state_dict(), weights_path) # Save the best threshold alongside weights so run_test.py can load it meta = {"threshold": best_threshold, "val_step_acc": best_val_step_acc, "val_f1": metrics["judgment_f1"], "epoch": epoch} with open(weights_path.replace(".pth", "_meta.json"), "w") as f: json.dump(meta, f, indent=2) print(f" [✓] New best model saved (step-acc={best_val_step_acc*100:.1f}%, " f"thr={best_threshold:.2f})") else: epochs_no_improve += 1 print(f" No improvement for {epochs_no_improve}/{args.patience} epochs.") if epochs_no_improve >= args.patience: print(f"\nEarly stopping at epoch {epoch}.") break print(f"\n── Training complete ──────────────────────────────────────────") print(f"Best val step-acc : {best_val_step_acc*100:.1f}% at threshold {best_threshold:.2f}") print(f"Weights saved to : {weights_path}") if __name__ == "__main__": main()