"""Qwen3-VL chat-template dataset for CoT + per-frame BeliefToken SFT. Two supervision modes (auto-detected per record): (1) Per-frame POMDP target — when belief.actions_per_frame is present: Scene: {scene} Critical: {critical} Threat: {threat} <|BELIEF|> <|A_0|> <|BELIEF|> <|A_1|> ... <|BELIEF|> <|A_{T-1}|> (2) Clip-level (legacy) — when only belief.action is present: Scene: {scene} Critical: {critical} Threat: {threat} <|BELIEF|> <|ACTION|> At SFT time only the assistant tokens receive gradient (prefix masked with -100). At belief-extraction time we teacher-force the full assistant string and read `last_hidden_state` at each `<|BELIEF|>` position — T 2560-D vectors per clip. """ from __future__ import annotations import json from pathlib import Path from typing import Any, Dict, List import torch from torch.utils.data import Dataset from training.VLA.frame_utils import sample_frames, sample_frames_from_mp4 # ───────────────────── special tokens ───────────────────── 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 = {"ALERT": 0, "OBSERVE": 1, "SILENT": 2} # ───────────────────── prompts ───────────────────── SYSTEM_PROMPT = ( "You are a driving-safety assistant. Given N dashcam frames (earliest → latest), " "produce a short chain-of-thought analysis, then emit one risk action token " "per frame wrapped in <|BELIEF|> ... . " "Actions: <|ALERT|> (collision < 0.5s), <|OBSERVE|> (threat 0.5-2.5s), " "<|SILENT|> (no threat). Keep prose minimal; the belief blocks are mandatory." ) USER_PROMPT = "Analyze the frames and emit scene analysis + per-frame belief blocks." def _parse_per_frame_belief(threat: str) -> Dict[int, str]: """Parse 'f0: phrase; f1: phrase; ...' into {frame_idx: phrase}.""" import re out = {} if not threat: return out parts = re.split(r"f(\d+):\s*", threat) # parts looks like ['', '0', 'phrase0;', '1', 'phrase1;', ...] for i in range(1, len(parts) - 1, 2): try: idx = int(parts[i]) phrase = parts[i + 1].strip().rstrip(";").strip() if phrase: out[idx] = phrase except (ValueError, IndexError): continue return out def _state_phrase_prefix(state: str) -> str: """Prefix that hints the model what kind of belief to encode per state. SILENT → broad scene context (lane / traffic / weather) OBSERVE → suspect agent + predicted trajectory ALERT → hazard itself + distance / urgency """ return { "SILENT": "context:", "OBSERVE": "watching:", "ALERT": "hazard:", }.get(state, "context:") def format_assistant_v4(beliefs_per_frame: List[str]) -> str: """v4 canonical assistant text: one <|BELIEF|> {scene+danger} per frame. No action token inside the span (action is emitted by the policy head downstream). This matches tools/make_cache_gt_belief.py. """ return "\n".join( f"{BELIEF_OPEN} {b.strip()} {BELIEF_CLOSE}" for b in beliefs_per_frame ) def format_assistant(cot: Dict[str, Any], actions: List[str], state_conditional: bool = False) -> str: """Build the exact assistant string the model must produce. `actions` is a list of action *names* (e.g. ["OBSERVE","OBSERVE","ALERT",...]). Single-element list degenerates to the legacy clip-level format. When `state_conditional=True`, emit per-frame state-specific phrases extracted from `cot.threat_analysis` *between* `<|BELIEF|>` and the action token (Stage A of VLAlert-X plan §B). The phrase content forces the BELIEF hidden state to encode different information per state. """ scene = str(cot.get("scene", "")).strip() critical = "; ".join(str(x).strip() for x in cot.get("critical_objects", []) if str(x).strip()) threat = str(cot.get("threat_analysis", "")).strip() lines = [f"Scene: {scene}", f"Critical: {critical}", f"Threat: {threat}"] if state_conditional: per_frame = _parse_per_frame_belief(threat) for i, a in enumerate(actions): phrase = per_frame.get(i, "").strip() # truncate phrase to ~15 words for token budget phrase = " ".join(phrase.split()[:15]) prefix = _state_phrase_prefix(a) if phrase: lines.append(f"{BELIEF_OPEN} {prefix} {phrase} " f"{ACTION_MAP[a]} {BELIEF_CLOSE}") else: # fallback to legacy if no per-frame phrase available lines.append(f"{BELIEF_OPEN} {ACTION_MAP[a]} {BELIEF_CLOSE}") else: for a in actions: lines.append(f"{BELIEF_OPEN} {ACTION_MAP[a]} {BELIEF_CLOSE}") return "\n".join(lines) def build_chat(frames, assistant_text: str | None): user_content = [{"type": "image", "image": img} for img in frames] user_content.append({"type": "text", "text": USER_PROMPT}) messages = [ {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, {"role": "user", "content": user_content}, ] if assistant_text is not None: messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_text}]}) return messages def _resolve_actions(belief: Dict[str, Any], n_frames: int) -> List[str]: """Prefer per-frame POMDP labels; fall back to clip-level repeated T times.""" pf = belief.get("actions_per_frame") if pf is not None and len(pf) > 0: if len(pf) < n_frames: pf = pf + [pf[-1]] * (n_frames - len(pf)) elif len(pf) > n_frames: pf = pf[:n_frames] return list(pf) return [belief["action"]] * n_frames class CoTBeliefDataset(Dataset): """Yields Qwen3-VL chat-template tensors with per-token labels. Requires the processor's tokenizer to ALREADY have the 5 special tokens added (via `add_special_tokens({"additional_special_tokens": ALL_SPECIAL})`). """ def __init__( self, jsonl_path: str, video_dir: str, processor, n_frames: int = 8, resize_short: int = 336, max_len: int = 4096, per_frame: bool = True, state_conditional: bool = False, video_root_override: str | None = None, ): self.video_dir = Path(video_dir) self.processor = processor self.n_frames = n_frames self.resize_short = resize_short self.max_len = max_len self.per_frame = per_frame self.state_conditional = state_conditional self.video_root_override = Path(video_root_override) if video_root_override else None self.records: List[Dict[str, Any]] = [] missing = 0 with open(jsonl_path) as f: for line in f: rec = json.loads(line) if rec.get("cot") is None or rec.get("belief") is None: missing += 1 continue self.records.append(rec) if missing: print(f"[CoTBeliefDataset] skipped {missing} records without cot+belief") def __len__(self): return len(self.records) def _resolve_video_path(self, rec: Dict[str, Any]) -> Path: if rec.get("video_path"): return Path(rec["video_path"]) clip_id = str(rec["id"]).zfill(5) return self.video_dir / f"{clip_id}.mp4" def __getitem__(self, idx): rec = self.records[idx] clip_id = str(rec["id"]) video_path = self._resolve_video_path(rec) frame_idx = rec.get("belief", {}).get("frame_indices") frames = sample_frames(video_path, n_frames=self.n_frames, resize_short=self.resize_short, frame_indices=frame_idx) if self.per_frame: actions = _resolve_actions(rec["belief"], self.n_frames) else: actions = [rec["belief"]["action"]] assistant_text = format_assistant(rec["cot"], actions, state_conditional=self.state_conditional) full_msgs = build_chat(frames, assistant_text=assistant_text) prefix_msgs = build_chat(frames, assistant_text=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_idx = [ACTION_TO_IDX[a] for a in actions] item = { "input_ids": input_ids, "attention_mask": full["attention_mask"][0], "labels": labels, "pixel_values": full["pixel_values"], "image_grid_thw": full["image_grid_thw"], "label": int(rec["label"]), "actions": actions, "action_idx": torch.tensor(action_idx, dtype=torch.long), "id": clip_id, } return item def collate_fn(batch, pad_token_id: int): max_len = max(b["input_ids"].size(0) for b in batch) input_ids, attn, labels, pixel_values, grid_thw = [], [], [], [], [] for b in batch: pad_n = max_len - b["input_ids"].size(0) input_ids.append(torch.cat([b["input_ids"], torch.full((pad_n,), pad_token_id, dtype=torch.long)])) attn.append(torch.cat([b["attention_mask"], torch.zeros(pad_n, dtype=b["attention_mask"].dtype)])) labels.append(torch.cat([b["labels"], torch.full((pad_n,), -100, dtype=torch.long)])) pixel_values.append(b["pixel_values"]) grid_thw.append(b["image_grid_thw"]) T = max(len(b["actions"]) for b in batch) action_idx = torch.full((len(batch), T), -1, dtype=torch.long) for i, b in enumerate(batch): action_idx[i, :len(b["actions"])] = b["action_idx"] return { "input_ids": torch.stack(input_ids), "attention_mask": torch.stack(attn), "labels": torch.stack(labels), "pixel_values": torch.cat(pixel_values, dim=0), "image_grid_thw": torch.cat(grid_thw, dim=0), "_clip_ids": [b["id"] for b in batch], "_actions": [b["actions"] for b in batch], "_action_idx": action_idx, "_cls_labels": torch.tensor([b["label"] for b in batch], dtype=torch.long), }