"""PyTorch Dataset that yields Qwen2.5-VL chat-template inputs for CoT SFT. Given a jsonl with {id, label, cot: {...}} records, for each sample we: 1. Extract 8 frames from the mp4. 2. Build a chat with system + user(images + prompt) + assistant(CoT JSON). 3. Tokenize with the processor; compute per-token label mask so that the LM loss is applied ONLY to the assistant's JSON tokens. The assistant message is forced to the exact JSON string (normalized order) so that the "verdict" key appears at a deterministic position — at inference time we can parse the "yes"/"no" token logit from a fixed slot. """ 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_from_mp4 SYSTEM_PROMPT = ( "You are a driving-safety assistant. Given 8 dashcam frames (earliest → latest), " "analyze the scene and output a JSON object with keys " "{scene, critical_objects, threat_analysis, verdict, confidence}. " "verdict is 'yes' if a collision or near-collision occurs in the clip, else 'no'. " "Output JSON only." ) USER_PROMPT = "Analyze the 8 frames and output the JSON." def canonical_cot_json(cot: Dict[str, Any]) -> str: """Re-serialize in a fixed key order so 'verdict' is always at the same place.""" ordered = { "scene": str(cot.get("scene", "")).strip(), "critical_objects": list(cot.get("critical_objects", []))[:4], "threat_analysis": str(cot.get("threat_analysis", "")).strip(), "verdict": str(cot.get("verdict", "no")).strip().lower(), "confidence": int(cot.get("confidence", 50)), } # compact separators (no spaces) so the downstream marker # '"verdict":"' matches exactly during inference. return json.dumps(ordered, ensure_ascii=False, separators=(",", ":")) 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 class NexarCoTDataset(Dataset): def __init__( self, jsonl_path: str, video_dir: str, processor, n_frames: int = 8, resize_short: int = 336, max_len: int = 4096, supervise: str = "assistant", # "assistant" | "verdict_only" ): 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.supervise = supervise assert supervise in ("assistant", "verdict_only"), supervise self.records: List[Dict[str, Any]] = [] with open(jsonl_path) as f: for line in f: rec = json.loads(line) if rec.get("cot") is None: continue self.records.append(rec) def __len__(self): return len(self.records) def __getitem__(self, idx): rec = self.records[idx] clip_id = str(rec["id"]).zfill(5) video_path = self.video_dir / f"{clip_id}.mp4" frames = sample_frames_from_mp4(video_path, n_frames=self.n_frames, resize_short=self.resize_short) assistant_text = canonical_cot_json(rec["cot"]) # Build the full chat (with assistant) and tokenize. 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 # mask system+user+prompt if self.supervise == "verdict_only": # Concentrate all gradient on the single yes/no token. # Strategy: find the LAST bare-token occurrence of "yes" (id 9693) or "no" (id 2152) # anywhere in the assistant region. Because canonical JSON puts "verdict" near the # end and bare yes/no (no leading space) are rare elsewhere, this reliably picks # the verdict slot without BPE boundary-alignment headaches. VERDICT_IDS = {9693, 2152} # "yes", "no" (Qwen2.5-VL tokenizer) assistant_ids = input_ids[prefix_len:].tolist() matches = [i + prefix_len for i, tid in enumerate(assistant_ids) if tid in VERDICT_IDS] labels[prefix_len:] = -100 assert matches, ( f"No verdict token (9693/2152) found in assistant region for id={rec['id']}. " f"Decoded assistant: {self.processor.tokenizer.decode(assistant_ids)[:400]}" ) labels[matches[-1]] = input_ids[matches[-1]] 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"]), "id": clip_id, } return item def collate_fn(batch, pad_token_id: int): """Right-pad variable-length sequences; stack images (Qwen2.5-VL already flattens).""" 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"]) 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], "_cls_labels": torch.tensor([b["label"] for b in batch], dtype=torch.long), }