training/VLA/cot_dataset.py

"""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),
    }