#!/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()