"""VLAlert-X v2 SFT dataset — per-frame BELIEF reasoning content. KEY DIFFERENCE from v1 (`cot_belief_dataset.py`): v1 wrote `<|BELIEF|> <|ACTION_i|> ` (action token wedged BETWEEN BELIEF tags → causes leak when pooling at BELIEF positions). v2 writes `<|BELIEF|> {per-frame reasoning text} <|ACTION_i|>` so BELIEF tags wrap actual REASONING and the action token sits AFTER the closing tag. Pooling inside the BELIEF span yields a leak-free perception vector; the action token never enters the pool window. Manifest schema expected (one record per tick, jsonl): { "id": str, "video_id": str, "video_path": str, "source": str, "frame_indices": [8 ints], "actions_per_frame": [8 strs of {SILENT, OBSERVE, ALERT}], "beliefs_per_frame": [8 strs, 10-25 tokens each], "danger_per_frame": [8 floats in [0, 1]], "tta_per_frame": [8 floats, seconds], "tick_action": str, "tick_tta_raw": float, "scene": str (optional, prepended if non-empty), "critical": str (optional, prepended if non-empty), ... } Assistant text format produced: [Scene: ...] ← optional [Critical: ...] ← optional <|BELIEF|> {belief_0} <|ACTION_0|> <|BELIEF|> {belief_1} <|ACTION_1|> ... <|BELIEF|> {belief_7} <|ACTION_7|> CE loss is on all assistant tokens (model must generate the belief text AND the action token). Belief content is teacher-forced from manifest during SFT so the model learns: visual → reasoning + action. For cache extraction (separate, see `tools/make_cache_x_v2.py`), action tokens are STRIPPED from the prompt so causal attention can't leak GT actions when we pool inside the BELIEF span. """ from __future__ import annotations import json from pathlib import Path from typing import Any, Dict, List, Optional import torch from torch.utils.data import Dataset from training.VLA.frame_utils import sample_frames # ───────────────────── special tokens (same as v1) ───────────────────── BELIEF_OPEN = "<|BELIEF|>" BELIEF_CLOSE = "" ACTION_ALERT = "<|ALERT|>" ACTION_OBSERVE = "<|OBSERVE|>" ACTION_SILENT = "<|SILENT|>" ACTION_TOKENS = [ACTION_ALERT, ACTION_OBSERVE, ACTION_SILENT] ALL_SPECIAL = [BELIEF_OPEN, BELIEF_CLOSE] + ACTION_TOKENS ACTION_MAP = { "ALERT": ACTION_ALERT, "OBSERVE": ACTION_OBSERVE, "SILENT": ACTION_SILENT, } ACTION_TO_IDX = {"SILENT": 0, "OBSERVE": 1, "ALERT": 2} # ───────────────────── prompts ───────────────────── SYSTEM_PROMPT_V2 = ( "You are a driving-safety assistant. Given N dashcam frames " "(earliest → latest), for each frame produce a short reasoning sentence " "describing the most safety-relevant cue you observe (lead-vehicle behaviour, " "TTC estimate, pedestrians, sudden brake, lane drift, etc.), wrap it in " "<|BELIEF|>..., then immediately emit the per-frame action: " "<|SILENT|> (no threat), <|OBSERVE|> (developing situation), " "or <|ALERT|> (imminent collision risk, < 2 s)." ) USER_PROMPT_V2 = ( "Emit 8 per-frame belief+action blocks for these frames." ) def format_assistant_v2(beliefs_per_frame: List[str], actions_per_frame: List[str], scene: str = "", critical: str = "") -> str: """Build the assistant string for v2 SFT. `beliefs_per_frame` must have length 8 (one per frame). `actions_per_frame` must have length 8, values in {SILENT, OBSERVE, ALERT}. `scene` and `critical` are optional clip-level prefix lines. """ assert len(beliefs_per_frame) == 8, "expected 8 belief sentences" assert len(actions_per_frame) == 8, "expected 8 actions" lines: List[str] = [] scene = (scene or "").strip() critical = (critical or "").strip() if scene: lines.append(f"Scene: {scene}") if critical: lines.append(f"Critical: {critical}") if lines: lines.append("") # blank line before frame blocks for b, a in zip(beliefs_per_frame, actions_per_frame): b_clean = (b or "").strip().replace("\n", " ") # cap at ~25 words to keep sequence length manageable b_clean = " ".join(b_clean.split()[:25]) action_tok = ACTION_MAP.get(a, ACTION_SILENT) lines.append(f"{BELIEF_OPEN} {b_clean} {BELIEF_CLOSE} {action_tok}") return "\n".join(lines) def build_chat_v2(frames, assistant_text: Optional[str]): user_content = [{"type": "image", "image": img} for img in frames] user_content.append({"type": "text", "text": USER_PROMPT_V2}) msgs = [ {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT_V2}]}, {"role": "user", "content": user_content}, ] if assistant_text is not None: msgs.append({"role": "assistant", "content": [{"type": "text", "text": assistant_text}]}) return msgs # ───────────────────── Dataset ───────────────────── class CoTBeliefDatasetV2(Dataset): """Per-frame BELIEF reasoning SFT dataset. Requires the processor's tokenizer to already have ALL_SPECIAL added. """ def __init__(self, jsonl_path: str, processor, n_frames: int = 8, resize_short: int = 336, max_len: int = 4096, action_token_weight: float = 2.0): """ action_token_weight: 2.0 → action token positions get 2× CE weight (encourages crisp action prediction; tracked via returned `action_token_mask`). """ self.processor = processor self.n_frames = n_frames self.resize_short = resize_short self.max_len = max_len self.action_token_weight = action_token_weight self.records: List[Dict[str, Any]] = [] n_skipped = 0 with open(jsonl_path) as f: for ln in f: ln = ln.strip() if not ln: continue try: r = json.loads(ln) except json.JSONDecodeError: continue # validate required fields ok = (isinstance(r.get("beliefs_per_frame"), list) and len(r["beliefs_per_frame"]) == n_frames and isinstance(r.get("actions_per_frame"), list) and len(r["actions_per_frame"]) == n_frames and isinstance(r.get("frame_indices"), list) and len(r["frame_indices"]) == n_frames and r.get("video_path")) if not ok: n_skipped += 1 continue self.records.append(r) print(f"[CoTBeliefDatasetV2] loaded {len(self.records)} records " f"(skipped {n_skipped} malformed) from {jsonl_path}") # cache action token ids for action_token_mask tok = processor.tokenizer self.action_ids = set() for t in ACTION_TOKENS: tid = tok.convert_tokens_to_ids(t) if tid is not None and tid != tok.unk_token_id: self.action_ids.add(tid) def __len__(self): return len(self.records) def __getitem__(self, idx): rec = self.records[idx] # sample frames frames = sample_frames( rec["video_path"], n_frames=self.n_frames, resize_short=self.resize_short, frame_indices=rec["frame_indices"], ) # build assistant text assistant_text = format_assistant_v2( beliefs_per_frame=rec["beliefs_per_frame"], actions_per_frame=rec["actions_per_frame"], scene=rec.get("scene", ""), critical=rec.get("critical", ""), ) full_msgs = build_chat_v2(frames, assistant_text) prefix_msgs = build_chat_v2(frames, None) proc = self.processor full_text = proc.apply_chat_template(full_msgs, tokenize=False, add_generation_prompt=False) prefix_text = proc.apply_chat_template(prefix_msgs, tokenize=False, add_generation_prompt=True) full = proc(text=[full_text], images=[frames], return_tensors="pt", padding=False, truncation=True, max_length=self.max_len) prefix = proc(text=[prefix_text], images=[frames], return_tensors="pt", padding=False, truncation=True, max_length=self.max_len) input_ids = full["input_ids"][0] labels = input_ids.clone() prefix_len = prefix["input_ids"].shape[1] labels[:prefix_len] = -100 # action token mask for weighted CE action_mask = torch.zeros_like(input_ids, dtype=torch.bool) for i, tid in enumerate(input_ids.tolist()): if i >= prefix_len and tid in self.action_ids: action_mask[i] = True # NOTE: pixel_values and image_grid_thw are kept unsliced (per-image # flat layout that Qwen3-VL processor returns) so the collator can # torch.cat across batch dim, matching v1 conventions. item = { "input_ids": input_ids, "labels": labels, "action_token_mask": action_mask, "attention_mask": full["attention_mask"][0] if "attention_mask" in full else None, "pixel_values": full["pixel_values"] if "pixel_values" in full else None, "image_grid_thw": full["image_grid_thw"] if "image_grid_thw" in full else None, } for k in ("video_grid_thw", "pixel_values_videos"): if k in full: item[k] = full[k] return item # ───────────────────── Collator ───────────────────── class CollatorV2: """Pad seq dim; cat pixel/grid along their natural dim (matches Qwen3-VL).""" def __init__(self, processor, n_frames: int = 8): self.processor = processor self.n_frames = n_frames self.pad_id = (processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id or 0) def __call__(self, batch): max_len = max(b["input_ids"].size(0) for b in batch) ids = torch.full((len(batch), max_len), self.pad_id, dtype=torch.long) labs = torch.full((len(batch), max_len), -100, dtype=torch.long) amask = torch.zeros((len(batch), max_len), dtype=torch.bool) attn_mask = torch.zeros((len(batch), max_len), dtype=torch.long) for i, b in enumerate(batch): L = b["input_ids"].size(0) ids[i, :L] = b["input_ids"] labs[i, :L] = b["labels"] amask[i, :L] = b["action_token_mask"] if b.get("attention_mask") is not None: attn_mask[i, :L] = b["attention_mask"] else: attn_mask[i, :L] = 1 out = { "input_ids": ids, "labels": labs, "attention_mask": attn_mask, "action_token_mask": amask, } # pixel_values: shape [num_patches_total, dim] — cat across batch if batch[0].get("pixel_values") is not None: out["pixel_values"] = torch.cat([b["pixel_values"] for b in batch], dim=0) # image_grid_thw: shape [n_images_per_sample, 3] — cat across batch if batch[0].get("image_grid_thw") is not None: out["image_grid_thw"] = torch.cat([b["image_grid_thw"] for b in batch], dim=0) for k in ("video_grid_thw", "pixel_values_videos"): if batch[0].get(k) is not None: out[k] = torch.cat([b[k] for b in batch], dim=0) return out