tools/score_vlalert_x_v2.py

"""VLAlert-X v2 Phase 5 — score a benchmark via cached features + heads.

Load a dual-stream cache (from `tools/make_cache_x_v2.py`) and the trained
DangerHead + PolicyHeadV2 checkpoints. Forward to produce per-tick action
probabilities, then save in the standard `per_tick.pt` schema that the
existing `tools/compute_daus_*.py` utilities consume.

Output schema:
  {
    "ids":           list[str]   (N,)
    "indices":       LongTensor  [N]
    "scores_binary": FloatTensor [N, 1]     # P(ALERT)
    "scores_3class": FloatTensor [N, 1, 3]  # P(S), P(O), P(A)
    "tta_per_tick":  FloatTensor [N, 1]
    "frame_indices": LongTensor  [N, 8]
    "category":      list[str]
    "source":        list[str]
    "tta_raw":       FloatTensor [N]
    "n_ticks":       int = 1
    "method":        "VLAlert-X-v2"
    "danger_ckpt":   str
    "policy_ckpt":   str
  }

Usage:
    python tools/score_vlalert_x_v2.py \
        --cache       data/belief_cache_v2/sft_x_v2__multisrc_val_full.pt \
        --manifest    data/cot_corpus_v2/multisrc_val_full_perframe_v2.jsonl \
        --danger_ckpt checkpoints/danger_v2/seed2/best.pt \
        --policy_ckpt checkpoints/policy_v2/seed2/best.pt \
        --out         eval_results/aus_metric/multisrc_per_tick/vlalert_x_v2.pt
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Dict

import torch
import torch.nn.functional as F
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from lkalert.models.danger_head import DangerHead
from lkalert.models.policy_head_v2 import PolicyHeadV2

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


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache",       type=Path, required=True,
                    help="Dual-stream cache .pt from tools/make_cache_x_v2.py")
    ap.add_argument("--manifest",    type=Path, required=True,
                    help="Perframe-v2 jsonl that was used to build the cache "
                         "(needed for category/source/tta_raw metadata)")
    ap.add_argument("--danger_ckpt", type=Path, required=True)
    ap.add_argument("--policy_ckpt", type=Path, required=True)
    ap.add_argument("--out",         type=Path, required=True)
    ap.add_argument("--batch_size",  type=int, default=64)
    ap.add_argument("--prev_action", type=int, default=3,
                    help="prev_action embedding index; 3=BOS (no temporal context)")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # ── load cache ──
    logger.info(f"[load] cache: {args.cache}")
    d = torch.load(args.cache, weights_only=False, map_location="cpu")
    belief  = d["belief_content"].float()        # [N, 8, D_belief]
    policy  = d["policy_position"].float()       # [N, 8, D_policy]
    valid   = d["valid_frames"]                  # [N, 8] bool
    ids_cache = list(d["ids"])
    N = belief.shape[0]
    logger.info(f"  N={N}  belief={tuple(belief.shape)}  policy={tuple(policy.shape)}")

    # ── load manifest for metadata (category, source, tta_raw, frame_indices) ──
    logger.info(f"[load] manifest: {args.manifest}")
    meta_by_id: Dict[str, Dict] = {}
    with open(args.manifest) as f:
        for ln in f:
            ln = ln.strip()
            if not ln: continue
            r = json.loads(ln)
            mid = r.get("id") or r.get("video_id")
            if mid:
                meta_by_id[mid] = r
    logger.info(f"  manifest records: {len(meta_by_id)}")

    # ── load heads ──
    logger.info(f"[load] DangerHead: {args.danger_ckpt}")
    ck_d = torch.load(args.danger_ckpt, weights_only=False, map_location="cpu")
    danger = DangerHead(in_dim=ck_d["in_dim"]).to(device)
    danger.load_state_dict(ck_d["model"])
    danger.eval()

    logger.info(f"[load] PolicyHeadV2: {args.policy_ckpt}")
    ck_p = torch.load(args.policy_ckpt, weights_only=False, map_location="cpu")
    policy_head = PolicyHeadV2(
        policy_dim=ck_p["policy_dim"],
        perception_dim_per_query=ck_p["perception_dim_per_query"],
        k_queries=ck_p["k_queries"],
    ).to(device)
    policy_head.load_state_dict(ck_p["model"])
    policy_head.eval()
    logger.info(f"  Phase 3 val: per_frame_auc={ck_d['val_metrics'].get('per_frame_auc',0):.4f}")
    logger.info(f"  Phase 4 val: bal_acc={ck_p['val_metrics']['balanced_acc']:.4f}  "
                f"per_class_recall={ck_p['val_metrics']['per_class_recall']}")

    # ── infer per-tick scores ──
    scores_3class = torch.zeros(N, 1, 3, dtype=torch.float32)
    n_failed = 0
    prev_act_tensor = torch.full((args.batch_size,), args.prev_action, dtype=torch.long, device=device)

    bs = args.batch_size
    for i in tqdm(range(0, N, bs), ncols=80, desc="infer"):
        end = min(N, i + bs)
        b_belief = belief[i:end].to(device, non_blocking=True)
        b_policy = policy[i:end].to(device, non_blocking=True)
        b_valid  = valid[i:end].to(device, non_blocking=True)
        cur_bs = end - i

        # Danger forward → perception_summary + danger_per_frame
        d_out = danger(b_belief, valid_frames=b_valid)
        perc  = d_out["perception_summary"]            # [B, K, hidden]
        danger_pf = d_out["per_frame"]                 # [B, 8]

        # Policy forward
        prev = prev_act_tensor[:cur_bs]
        logits = policy_head(b_policy, perc, danger_pf, prev, valid_frames=b_valid)  # [B, 3]
        probs = F.softmax(logits, dim=-1).cpu()
        scores_3class[i:end, 0] = probs

    scores_binary = scores_3class[:, :, 2].clone()   # P(ALERT)

    # ── assemble per_tick.pt metadata ──
    ids_out: list = []
    cat_out: list = []
    src_out: list = []
    tta_raw_out = torch.zeros(N, dtype=torch.float32)
    tta_per_tick_out = torch.zeros(N, 1, dtype=torch.float32)
    frame_indices_out = torch.zeros(N, 8, dtype=torch.long)
    indices_out = torch.arange(N, dtype=torch.long)

    # IMPORTANT: cache stores per-tick category/source already (correctly
    # tied to each specific tick's TTA). Manifest meta_by_id dedups on
    # video_id and clobbers earlier ticks' category — DON'T use it for
    # category/source. Only use manifest for `frame_indices` lookup.
    cache_category = list(d.get("category", [""] * N))
    cache_source   = list(d.get("source",   [""] * N))
    cache_tick_tta = d.get("tick_tta_raw", torch.full((N,), -1.0))

    n_missing_meta = 0
    for i, vid in enumerate(ids_cache):
        m = meta_by_id.get(vid, {})
        if not m:
            n_missing_meta += 1
        ids_out.append(vid)
        cat_out.append(cache_category[i] if i < len(cache_category) else "")
        src_out.append(cache_source[i]   if i < len(cache_source)   else "")
        tta_v = (cache_tick_tta[i].item() if hasattr(cache_tick_tta[i], "item")
                  else float(cache_tick_tta[i]))
        tta_raw_out[i] = tta_v
        tta_per_tick_out[i, 0] = tta_v
        fi = m.get("frame_indices", [0]*8)
        frame_indices_out[i] = torch.tensor(fi[:8], dtype=torch.long)
    if n_missing_meta:
        logger.warning(f"  {n_missing_meta} cache ids had no matching manifest record "
                       f"(only frame_indices lost; category/source still correct from cache)")

    out_dict = {
        "ids":           ids_out,
        "indices":       indices_out,
        "scores_binary": scores_binary,
        "scores_3class": scores_3class,
        "tta_per_tick":  tta_per_tick_out,
        "frame_indices": frame_indices_out,
        "category":      cat_out,
        "source":        src_out,
        "tta_raw":       tta_raw_out,
        "n_ticks":       1,
        "method":        "VLAlert-X-v2",
        "danger_ckpt":   str(args.danger_ckpt),
        "policy_ckpt":   str(args.policy_ckpt),
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    torch.save(out_dict, args.out)
    logger.info(f"[save] {args.out}")
    logger.info(f"  N={len(ids_out)}  "
                f"P(ALERT) range=[{scores_binary.min():.4f}, {scores_binary.max():.4f}]")
    # distribution by category
    from collections import Counter
    cc = Counter(cat_out)
    logger.info(f"  category dist: {dict(cc)}")


if __name__ == "__main__":
    main()