training/Nexar/nexar_trainer.py

#!/usr/bin/env python3
"""
Train a NexarTemporalHead (or NexarSimpleHead) on domain-adapted Nexar features.

The SFT backbone is already frozen and features are pre-cached. This trainer
only optimises the lightweight collision prediction head.

Usage:
  python -m training.Nexar.nexar_trainer \
      --cache_pos  data/nexar_cache/train_positive.pt \
      --cache_neg  data/nexar_cache/train_negative.pt \
      --output_dir checkpoints/Nexar/nexar_v1 \
      --arch       temporal \
      --n_windows  3 \
      --epochs     20 \
      --batch_size 64 \
      --lr         3e-4
"""
from __future__ import annotations

import argparse
import json
import logging
import random
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader, WeightedRandomSampler
from sklearn.metrics import average_precision_score, roc_auc_score

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from training.Nexar.nexar_dataset import NexarTrainDataset, nexar_collate_train
from training.Nexar.nexar_model import build_model

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Nexar.trainer")

SEED = 42


def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def split_dataset(ds: NexarTrainDataset, val_frac: float = 0.15, seed: int = SEED):
    """Stratified train/val split."""
    labels = [s["label"] for s in ds.samples]
    pos_idx = [i for i, l in enumerate(labels) if l == 1]
    neg_idx = [i for i, l in enumerate(labels) if l == 0]

    rng = random.Random(seed)
    rng.shuffle(pos_idx)
    rng.shuffle(neg_idx)

    n_val_pos = max(1, int(len(pos_idx) * val_frac))
    n_val_neg = max(1, int(len(neg_idx) * val_frac))

    val_idx  = pos_idx[:n_val_pos] + neg_idx[:n_val_neg]
    train_idx = pos_idx[n_val_pos:] + neg_idx[n_val_neg:]

    from torch.utils.data import Subset
    return Subset(ds, train_idx), Subset(ds, val_idx)


def make_sampler(subset) -> WeightedRandomSampler:
    """Class-balanced weighted sampler for the training subset."""
    labels = [subset.dataset.samples[i]["label"] for i in subset.indices]
    labels_arr = np.array(labels, dtype=float)
    n_pos = labels_arr.sum()
    n_neg = len(labels_arr) - n_pos
    weights = np.where(labels_arr == 1, len(labels_arr) / (2 * n_pos + 1e-9),
                                         len(labels_arr) / (2 * n_neg + 1e-9))
    return WeightedRandomSampler(
        weights=torch.from_numpy(weights).float(),
        num_samples=len(subset),
        replacement=True,
    )


def compute_ap(labels: np.ndarray, scores: np.ndarray) -> float:
    try:
        return float(average_precision_score(labels, scores))
    except Exception:
        return float("nan")


def train_epoch(model, loader, optimizer, device) -> float:
    model.train()
    total_loss = 0.0
    n = 0
    for batch in loader:
        beliefs   = batch["beliefs"].to(device)    # [B, T, H] or [B, H]
        tta_means = batch["tta_means"].to(device)
        tta_vars  = batch["tta_vars"].to(device)
        p_alerts  = batch["p_alerts"].to(device)
        labels    = batch["labels"].to(device)     # [B] float

        if isinstance(model, torch.nn.Module) and hasattr(model, "lstm"):
            # Temporal model: [B, T, H]
            scores = model(beliefs, tta_means, tta_vars, p_alerts)
        else:
            # Simple model: last window only [B, H]
            scores = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1])

        loss = F.binary_cross_entropy(scores, labels)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        total_loss += loss.item() * len(labels)
        n += len(labels)
    return total_loss / max(n, 1)


@torch.no_grad()
def eval_epoch(model, loader, device) -> Tuple[float, float, float]:
    model.eval()
    all_scores: List[float] = []
    all_labels: List[float] = []
    total_loss = 0.0
    n = 0
    for batch in loader:
        beliefs   = batch["beliefs"].to(device)
        tta_means = batch["tta_means"].to(device)
        tta_vars  = batch["tta_vars"].to(device)
        p_alerts  = batch["p_alerts"].to(device)
        labels    = batch["labels"].to(device)

        if hasattr(model, "lstm"):
            scores = model(beliefs, tta_means, tta_vars, p_alerts)
        else:
            scores = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1])

        loss = F.binary_cross_entropy(scores, labels)
        total_loss += loss.item() * len(labels)
        n += len(labels)

        all_scores.extend(scores.cpu().tolist())
        all_labels.extend(labels.cpu().tolist())

    arr_l = np.array(all_labels)
    arr_s = np.array(all_scores)
    ap  = compute_ap(arr_l, arr_s)
    try:
        auc = float(roc_auc_score(arr_l, arr_s))
    except Exception:
        auc = float("nan")
    return total_loss / max(n, 1), ap, auc


def main():
    parser = argparse.ArgumentParser("nexar_trainer")
    parser.add_argument("--cache_pos",   required=True,  help=".pt cache for positive train videos")
    parser.add_argument("--cache_neg",   required=True,  help=".pt cache for negative train videos")
    parser.add_argument("--output_dir",  required=True)
    parser.add_argument("--arch",        default="temporal", choices=["simple", "temporal"])
    parser.add_argument("--n_windows",   type=int,   default=3)
    parser.add_argument("--epochs",      type=int,   default=30)
    parser.add_argument("--batch_size",  type=int,   default=64)
    parser.add_argument("--lr",          type=float, default=3e-4)
    parser.add_argument("--lr_min",      type=float, default=1e-6)
    parser.add_argument("--weight_decay",type=float, default=1e-4)
    parser.add_argument("--val_frac",    type=float, default=0.15)
    parser.add_argument("--patience",    type=int,   default=8)
    parser.add_argument("--hidden_dim",  type=int,   default=2048,
                        help="SFT hidden_dim (Qwen2.5-VL-3B = 2048)")
    args = parser.parse_args()

    set_seed(SEED)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ── data ─────────────────────────────────────────────────────────────────
    full_ds = NexarTrainDataset(args.cache_pos, args.cache_neg, n_windows=args.n_windows)
    train_subset, val_subset = split_dataset(full_ds, val_frac=args.val_frac)
    logger.info(f"Train: {len(train_subset)}  Val: {len(val_subset)}")

    sampler = make_sampler(train_subset)
    train_loader = DataLoader(train_subset, batch_size=args.batch_size,
                              sampler=sampler, num_workers=4, collate_fn=nexar_collate_train,
                              pin_memory=True)
    val_loader   = DataLoader(val_subset,   batch_size=args.batch_size,
                              shuffle=False, num_workers=4, collate_fn=nexar_collate_train,
                              pin_memory=True)

    # ── model ─────────────────────────────────────────────────────────────────
    model = build_model(args.hidden_dim, args.arch).to(device)
    total_params = sum(p.numel() for p in model.parameters())
    logger.info(f"NexarHead ({args.arch}): {total_params:,} params")

    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    total_steps = args.epochs * len(train_loader)
    scheduler = CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=args.lr_min)

    # ── training loop ─────────────────────────────────────────────────────────
    best_ap   = 0.0
    patience_count = 0
    history = []

    for epoch in range(1, args.epochs + 1):
        train_loss = train_epoch(model, train_loader, optimizer, device)
        scheduler.step()
        val_loss, val_ap, val_auc = eval_epoch(model, val_loader, device)
        lr = optimizer.param_groups[0]["lr"]

        logger.info(
            f"Epoch {epoch:3d}/{args.epochs}  "
            f"train_loss={train_loss:.4f}  val_loss={val_loss:.4f}  "
            f"val_AP={val_ap:.4f}  val_AUC={val_auc:.4f}  lr={lr:.2e}"
        )
        history.append({
            "epoch": epoch, "train_loss": train_loss,
            "val_loss": val_loss, "val_ap": val_ap, "val_auc": val_auc,
        })

        if val_ap > best_ap:
            best_ap = val_ap
            patience_count = 0
            torch.save(model.state_dict(), out_dir / "best_model.pt")
            with open(out_dir / "best_meta.json", "w") as f:
                json.dump({"epoch": epoch, "val_ap": val_ap, "val_auc": val_auc,
                           "arch": args.arch, "hidden_dim": args.hidden_dim,
                           "n_windows": args.n_windows}, f, indent=2)
            logger.info(f"  ★ New best val_AP={best_ap:.4f} — checkpoint saved")
        else:
            patience_count += 1
            if patience_count >= args.patience:
                logger.info(f"Early stopping at epoch {epoch} (patience={args.patience})")
                break

    with open(out_dir / "history.json", "w") as f:
        json.dump(history, f, indent=2)

    logger.info(f"\n✅ Training complete. Best val_AP = {best_ap:.4f}")
    logger.info(f"   Checkpoint: {out_dir}/best_model.pt")


if __name__ == "__main__":
    main()