training/Policy/make_nexar_belief_cache.py

#!/usr/bin/env python3
"""
make_nexar_belief_cache.py
═══════════════════════════════════════════════════════════════════════════════
Nexar-only per-frame belief cache extractor for the CoT+BeliefToken
Qwen3-VL-4B checkpoint.

Why separate from make_cot_belief_cache.py:
  make_cot_belief_cache.py is bound to PolicyDataset, which expects
  pre-computed frame_indices per sample. The Nexar pipeline works straight
  from mp4s (Kaggle train.csv / sample_submission.csv), so we sample frames
  on the fly with training.VLA.frame_utils.sample_frames_from_mp4 and reuse
  the model-loading + token-splitting helpers from make_cot_belief_cache.

Input manifest (produced by tools/make_nexar_mp4_manifest.py):
  {
    "samples": [
      {"video_id": "00001", "mp4": "...", "label": 0 | 1 | -1,
       "time_of_alert": float | None, "time_of_event": float | None},
      ...
    ]
  }

Output (.pt):
  beliefs_frame [N, T, D] fp16  — per-frame pooled visual token hiddens
  valid_frames  [N, T]    bool  — True where frame was present
  beliefs_text  [N, D]    fp16  — mean of non-image valid tokens
  labels        [N]       int64 — 0 safe / 1 collision / -1 unknown (test)
  meta          dict

Usage
─────
  # Val (held-out ~20% of Kaggle train)
  python -m training.Policy.make_nexar_belief_cache --manifest data/nexar_mp4_manifest/val.json   --out data/belief_cache_nexar_qwen3vl4b/val.pt   --ckpt_dir checkpoints/VLA/qwen3vl4b_cot_belief/best --n_frames 8

  # Train
  python -m training.Policy.make_nexar_belief_cache --manifest data/nexar_mp4_manifest/train.json --out data/belief_cache_nexar_qwen3vl4b/train.pt --ckpt_dir checkpoints/VLA/qwen3vl4b_cot_belief/best --n_frames 8

  # Test (for submission)
  python -m training.Policy.make_nexar_belief_cache --manifest data/nexar_mp4_manifest/test.json  --out data/belief_cache_nexar_qwen3vl4b/test.pt  --ckpt_dir checkpoints/VLA/qwen3vl4b_cot_belief/best --n_frames 8
"""
from __future__ import annotations

import argparse
import json
import logging
import shutil
import sys
from pathlib import Path
from typing import Dict, List, Optional

import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from training.Policy.make_cot_belief_cache import (
    SYSTEM_PROMPT, USER_PROMPT,
    _config_hidden_size, _config_spatial_merge_size,
    _resize_short, extract_batch, load_model,
)
from training.VLA.frame_utils import (
    sample_frames_from_mp4, sample_frames_from_mp4_by_indices,
)
from training.Policy.policy_dataset import _resample_indices, SAMPLING_SCHEMES

import cv2  # for total-frame probe when --sampling != uniform

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


class NexarMp4Dataset(Dataset):
    """Reads an mp4 manifest and returns PIL frames + metadata per item."""

    def __init__(self, manifest_path: str | Path, n_frames: int = 8,
                 resize_short: int = 336, sampling: str = "uniform",
                 anchor_offset_seconds: float = 0.0):
        """anchor_offset_seconds shifts the SAMPLING WINDOW end backward by N
        seconds (e.g., 1.0 means sample from [0, clip_end - 1.0s]). This is
        the per-frame sliding-inference building block: extract one cache
        per offset, then aggregate."""
        path = Path(manifest_path)
        with open(path) as f:
            payload = json.load(f)
        self.samples: List[dict] = payload["samples"] if isinstance(payload, dict) else payload
        self.n_frames = n_frames
        self.resize_short = resize_short
        if sampling not in SAMPLING_SCHEMES:
            raise ValueError(f"unknown sampling: {sampling}")
        self.sampling = sampling
        self.anchor_offset_seconds = float(anchor_offset_seconds)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx: int):
        s = self.samples[idx]
        if self.sampling == "uniform" and self.anchor_offset_seconds == 0:
            frames = sample_frames_from_mp4(s["mp4"], self.n_frames,
                                            self.resize_short)
        else:
            cap = cv2.VideoCapture(str(s["mp4"]))
            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
            cap.release()
            if total <= 0:
                raise RuntimeError(f"bad video: {s['mp4']}")
            # Shift window end back by anchor_offset_seconds
            offset_frames = int(round(self.anchor_offset_seconds * fps))
            end = max(self.n_frames - 1, total - 1 - offset_frames)
            base = list(range(end + 1))
            indices = _resample_indices(base, self.n_frames, self.sampling)
            frames = sample_frames_from_mp4_by_indices(
                s["mp4"], indices, resize_short=self.resize_short,
            )
        return {
            "video_id": s["video_id"],
            "label":    int(s.get("label", -1)),
            "frames":   frames,
            "toa":      s.get("time_of_alert"),
            "toe":      s.get("time_of_event"),
        }


def _collate(batch):
    return {
        "video_ids": [b["video_id"] for b in batch],
        "labels":    [b["label"]    for b in batch],
        "frames":    [b["frames"]   for b in batch],
        "toa":       [b["toa"]      for b in batch],
        "toe":       [b["toe"]      for b in batch],
    }


def _build_inputs(processor, images_b: List[List], resize_short: int):
    """Match the chat template used during CoT+BeliefToken training (no assistant turn)."""
    images_b_resized = [[_resize_short(img, resize_short) for img in frames]
                        for frames in images_b]
    texts: List[str] = []
    for frames in images_b_resized:
        user_content = [{"type": "image", "image": img} for img in frames]
        user_content.append({"type": "text", "text": USER_PROMPT})
        msgs = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
            {"role": "user",   "content": user_content},
        ]
        texts.append(
            processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        )
    return processor(text=texts, images=images_b_resized,
                     return_tensors="pt", padding=True, truncation=False)


# ── chunked save/resume ────────────────────────────────────────────────────

def _flush_chunk(acc: Dict[str, List[torch.Tensor]], chunk_dir: Path, idx: int) -> int:
    if not acc:
        return 0
    part = {k: torch.cat(v, dim=0) for k, v in acc.items()}
    n = next(iter(part.values())).shape[0]
    tmp = chunk_dir / f"chunk_{idx:05d}.pt.tmp"
    fin = chunk_dir / f"chunk_{idx:05d}.pt"
    torch.save(part, tmp); tmp.rename(fin)
    return n


def _scan_chunks(chunk_dir: Path) -> int:
    if not chunk_dir.exists():
        return 0
    for t in chunk_dir.glob("*.tmp"):
        t.unlink(missing_ok=True)
    files = sorted(chunk_dir.glob("chunk_*.pt"))
    good = 0
    for f in files:
        try:
            torch.load(f, map_location="cpu", weights_only=True)
            good += 1
        except Exception as e:
            logger.warning(f"  [resume] dropping unreadable chunk {f.name}: {e}")
            f.unlink(missing_ok=True)
    return good


def _merge_chunks(chunk_dir: Path) -> Dict[str, torch.Tensor]:
    files = sorted(chunk_dir.glob("chunk_*.pt"))
    acc: Dict[str, List[torch.Tensor]] = {}
    for f in files:
        d = torch.load(f, map_location="cpu", weights_only=True)
        for k, v in d.items():
            acc.setdefault(k, []).append(v)
    return {k: torch.cat(lst, dim=0) for k, lst in acc.items()}


# ── main loop ──────────────────────────────────────────────────────────────

def build_cache(model, processor, loader: DataLoader,
                image_token_id: int, sms: int, n_frames: int,
                chunk_dir: Optional[Path], chunk_size: int,
                resize_short: int) -> Dict[str, torch.Tensor]:
    start_batch = 0
    chunk_idx = 0
    if chunk_dir is not None:
        chunk_dir.mkdir(parents=True, exist_ok=True)
        n_chunks = _scan_chunks(chunk_dir)
        if n_chunks > 0:
            start_batch = n_chunks * chunk_size
            chunk_idx = n_chunks
            logger.info(f"  [resume] {n_chunks} chunks; skipping first {start_batch} batches")

    acc: Dict[str, List[torch.Tensor]] = {}
    since_flush = 0
    pbar = tqdm(loader, desc="nexar-cache", ncols=80, leave=True)
    for bi, batch in enumerate(pbar):
        if bi < start_batch:
            continue
        inputs = _build_inputs(processor, batch["frames"], resize_short)
        feats = extract_batch(model, processor, inputs,
                              image_token_id, sms, n_frames)
        B = feats["beliefs_frame"].shape[0]
        feats["labels"]    = torch.tensor(batch["labels"], dtype=torch.long)
        # Keep video_ids per chunk for alignment on resume
        feats["video_idx"] = torch.tensor([bi * B + j for j in range(B)],
                                          dtype=torch.long)
        for k, v in feats.items():
            acc.setdefault(k, []).append(v)
        since_flush += 1
        if chunk_dir is not None and since_flush >= chunk_size:
            n = _flush_chunk(acc, chunk_dir, chunk_idx)
            pbar.set_postfix_str(f"chunk={chunk_idx} +{n}")
            acc = {}; since_flush = 0; chunk_idx += 1

    if chunk_dir is not None and acc:
        _flush_chunk(acc, chunk_dir, chunk_idx); acc = {}

    cache = _merge_chunks(chunk_dir) if chunk_dir is not None \
            else {k: torch.cat(lst, dim=0) for k, lst in acc.items()}
    return cache


def main():
    ap = argparse.ArgumentParser("make_nexar_belief_cache")
    ap.add_argument("--manifest",   required=True,
                    help="data/nexar_mp4_manifest/{train,val,test}.json")
    ap.add_argument("--out",        required=True)
    ap.add_argument("--ckpt_dir",   required=True,
                    help="PEFT adapter dir (CoT+BeliefToken checkpoint)")
    ap.add_argument("--base_model",
                    default="PROJECT_ROOT/models/Qwen3-VL-4B-Instruct")
    ap.add_argument("--sampling",   default="last_biased",
                    choices=list(SAMPLING_SCHEMES),
                    help="match make_cot_belief_cache (training uses last_biased)")
    ap.add_argument("--anchor_offset_seconds", type=float, default=0.0,
                    help="Shift sampling window end back by N seconds. "
                         "Use 0.5/1.0/1.5 to produce anchor variants for "
                         "per-frame sliding inference (Kaggle mAP boost).")
    ap.add_argument("--n_frames",   type=int, default=8,
                    help="Match training (CoT SFT used n_frames=8)")
    ap.add_argument("--resize_short", type=int, default=336)
    ap.add_argument("--batch_size", type=int, default=2,
                    help="VLM forward batch (2 doubles GPU util on 5090; >4 risks OOM)")
    ap.add_argument("--num_workers", type=int, default=8,
                    help="mp4 decode is the bottleneck — 6-8 workers saturate GPU")
    ap.add_argument("--prefetch_factor", type=int, default=4,
                    help="how many batches each worker pre-decodes")
    ap.add_argument("--chunk_size", type=int, default=200)
    ap.add_argument("--keep_chunks", action="store_true")
    ap.add_argument("--overwrite",  action="store_true")
    args = ap.parse_args()

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    if out_path.exists() and not args.overwrite:
        logger.info(f"Cache exists: {out_path} — use --overwrite to rebuild")
        return

    ds = NexarMp4Dataset(args.manifest, n_frames=args.n_frames,
                         resize_short=args.resize_short,
                         sampling=args.sampling,
                         anchor_offset_seconds=args.anchor_offset_seconds)
    if len(ds) == 0:
        raise SystemExit("manifest empty")
    logger.info(f"  {len(ds)} clips from {args.manifest}")

    model, processor = load_model(args.base_model, args.ckpt_dir)
    img_tok_id = processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")
    sms = _config_spatial_merge_size(model.config)
    hidden_dim = _config_hidden_size(model.config)
    logger.info(f"  image_token_id={img_tok_id} sms={sms} hidden_dim={hidden_dim}")

    loader = DataLoader(
        ds, batch_size=args.batch_size, shuffle=False,
        num_workers=args.num_workers, collate_fn=_collate,
        pin_memory=False,
        prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
        persistent_workers=args.num_workers > 0,
    )

    chunk_dir = out_path.parent / (out_path.stem + ".chunks") if args.chunk_size > 0 else None
    cache = build_cache(
        model, processor, loader,
        image_token_id=img_tok_id, sms=sms, n_frames=args.n_frames,
        chunk_dir=chunk_dir, chunk_size=args.chunk_size,
        resize_short=args.resize_short,
    )

    # Align video_ids list with samples order (guaranteed since shuffle=False)
    video_ids = [s["video_id"] for s in ds.samples]
    toas = [s.get("time_of_alert") for s in ds.samples]
    toes = [s.get("time_of_event") for s in ds.samples]

    n_out = int(cache["beliefs_frame"].shape[0])
    if n_out != len(video_ids):
        logger.warning(f"  n_cached={n_out} != n_manifest={len(video_ids)} "
                       f"(chunked resume may have dropped trailing batch; "
                       f"re-run with --overwrite if mismatch is large)")
        # Trim metadata lists to cache length to keep alignment
        video_ids = video_ids[:n_out]
        toas      = toas[:n_out]
        toes      = toes[:n_out]

    meta = {
        "schema_version":     3,
        "cache_mode":         "per_frame_cot_belief_nexar",
        "backbone":           "Qwen3-VL-4B-Instruct",
        "hidden_dim":         hidden_dim,
        "n_frames":           args.n_frames,
        "resize_short":       args.resize_short,
        "n_samples":          n_out,
        "spatial_merge_size": sms,
        "image_token_id":     int(img_tok_id),
        "ckpt_dir":           str(args.ckpt_dir),
        "base_model":         str(args.base_model),
        "manifest":           str(args.manifest),
        "video_ids":          video_ids,
        "time_of_alert":      toas,
        "time_of_event":      toes,
    }
    # drop helper tensor
    cache.pop("video_idx", None)

    to_save = dict(cache)
    to_save["meta"] = meta

    tmp = out_path.with_suffix(out_path.suffix + ".tmp")
    torch.save(to_save, tmp); tmp.rename(out_path)
    logger.info(f"  Saved -> {out_path}  ({n_out} clips)")

    with open(out_path.with_suffix(".meta.json"), "w") as f:
        slim = {k: v for k, v in meta.items()
                if k not in ("video_ids", "time_of_alert", "time_of_event")}
        slim["n_ids"] = len(video_ids)
        json.dump(slim, f, indent=2)

    if chunk_dir is not None and chunk_dir.exists() and not args.keep_chunks:
        shutil.rmtree(chunk_dir)
        logger.info(f"  removed {chunk_dir}")

    logger.info("done.")


if __name__ == "__main__":
    main()