training/VLA/cot_belief_dataset_v2.py

"""VLAlert-X v2 SFT dataset — per-frame BELIEF reasoning content.

KEY DIFFERENCE from v1 (`cot_belief_dataset.py`):
  v1 wrote `<|BELIEF|> <|ACTION_i|> </|BELIEF|>` (action token wedged BETWEEN
  BELIEF tags → causes leak when pooling at BELIEF positions).
  v2 writes `<|BELIEF|> {per-frame reasoning text} </|BELIEF|> <|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} </|BELIEF|> <|ACTION_0|>
  <|BELIEF|> {belief_1} </|BELIEF|> <|ACTION_1|>
  ...
  <|BELIEF|> {belief_7} </|BELIEF|> <|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   = "</|BELIEF|>"
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|>...</|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