trigger-off/src/training/data_collator.py

"""
DataCollator for OS-Atlas (InternVL2) training.

Handles batching of image + text pairs using InternVL2's conversation format:

    <|im_start|>user
    <image>
    INSTRUCTION<|im_end|>
    <|im_start|>assistant
    (x,y)<|im_end|>
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Sequence

import torch
from torch.utils.data import Dataset

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# InternVL2 conversation template constants
# ---------------------------------------------------------------------------

IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
IMAGE_TOKEN = "<image>"

USER_PREFIX = f"{IM_START}user\n{IMAGE_TOKEN}\n"
ASST_PREFIX = f"{IM_END}\n{IM_START}assistant\n"
ASST_SUFFIX = f"{IM_END}"


def build_internvl_prompt(instruction: str) -> str:
    """Build a complete InternVL2-format prompt string (without the response)."""
    return f"{USER_PREFIX}{instruction}{ASST_PREFIX}"


def build_internvl_full(instruction: str, response: str) -> str:
    """Build a complete InternVL2-format sequence including the response."""
    return f"{USER_PREFIX}{instruction}{ASST_PREFIX}{response}{ASST_SUFFIX}"


# ---------------------------------------------------------------------------
# OSAtlasDataCollator
# ---------------------------------------------------------------------------

class OSAtlasDataCollator:
    """
    Data collator for OS-Atlas fine-tuning.

    Takes a batch of dicts with keys:
      - "input_ids"    : (seq_len,) LongTensor
      - "labels"       : (seq_len,) LongTensor (masked prompt tokens → -100)
      - "pixel_values" : (C, H, W) FloatTensor
      - "attention_mask": (seq_len,) LongTensor  [optional, computed if absent]

    Returns a batched dict suitable for model.forward().
    """

    def __init__(
        self,
        pad_token_id: int = 0,
        label_pad_token_id: int = -100,
    ) -> None:
        self.pad_token_id = pad_token_id
        self.label_pad_token_id = label_pad_token_id

    def __call__(
        self, features: List[Dict[str, Any]]
    ) -> Dict[str, torch.Tensor]:
        """Collate a list of feature dicts into a batched tensor dict."""
        if not features:
            return {}

        batch: Dict[str, Any] = {}

        # ---- input_ids / labels ----
        input_ids_list = [f["input_ids"] for f in features if "input_ids" in f]
        labels_list = [f["labels"] for f in features if "labels" in f]

        if input_ids_list:
            batch["input_ids"] = self._pad_sequence(
                input_ids_list, pad_value=self.pad_token_id
            )
            batch["attention_mask"] = (batch["input_ids"] != self.pad_token_id).long()

        if labels_list:
            batch["labels"] = self._pad_sequence(
                labels_list, pad_value=self.label_pad_token_id
            )

        # ---- pixel_values ----
        pixel_values_list = [f["pixel_values"] for f in features if "pixel_values" in f]
        if pixel_values_list:
            try:
                batch["pixel_values"] = torch.stack(pixel_values_list, dim=0)
            except RuntimeError:
                # If shapes differ, pad to largest
                max_c = max(pv.shape[0] for pv in pixel_values_list)
                max_h = max(pv.shape[1] for pv in pixel_values_list)
                max_w = max(pv.shape[2] for pv in pixel_values_list)
                padded = []
                for pv in pixel_values_list:
                    p = torch.zeros(max_c, max_h, max_w, dtype=pv.dtype)
                    c, h, w = pv.shape
                    p[:c, :h, :w] = pv
                    padded.append(p)
                batch["pixel_values"] = torch.stack(padded, dim=0)

        # Pass through any other keys unchanged (e.g. image_flags)
        extra_keys = set(features[0].keys()) - {"input_ids", "labels", "pixel_values", "attention_mask"}
        for key in extra_keys:
            vals = [f[key] for f in features if key in f]
            if vals and isinstance(vals[0], torch.Tensor):
                try:
                    batch[key] = torch.stack(vals, dim=0)
                except RuntimeError:
                    batch[key] = vals
            else:
                batch[key] = vals

        return batch

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _pad_sequence(
        self,
        sequences: List[torch.Tensor],
        pad_value: int,
    ) -> torch.Tensor:
        """Right-pad a list of 1-D tensors to the same length."""
        max_len = max(s.shape[0] for s in sequences)
        out = torch.full(
            (len(sequences), max_len),
            fill_value=pad_value,
            dtype=sequences[0].dtype,
        )
        for i, s in enumerate(sequences):
            out[i, : s.shape[0]] = s
        return out


# ---------------------------------------------------------------------------
# SeeClick-compatible dataset wrapper
# ---------------------------------------------------------------------------

class SeeClickJSONDataset(Dataset):
    """
    Torch Dataset that loads SeeClick-format JSON conversation data.

    Compatible with both SeeClick and OS-Atlas collators.
    """

    def __init__(
        self,
        json_path: str,
        tokenizer: Any,
        image_processor: Any,
        max_length: int = 2048,
        model_type: str = "os_atlas",
    ) -> None:
        import json

        self.tokenizer = tokenizer
        self.image_processor = image_processor
        self.max_length = max_length
        self.model_type = model_type

        with open(json_path, "r", encoding="utf-8") as f:
            self.samples = json.load(f)

        logger.info("Loaded %d samples from %s.", len(self.samples), json_path)

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        sample = self.samples[idx]
        convs = sample["conversations"]

        user_text = ""
        asst_text = ""
        img_path = ""

        import re as _re

        for c in convs:
            if c["from"] == "user":
                value = c["value"]
                m = _re.search(r"<img>(.*?)</img>", value)
                if m:
                    img_path = m.group(1)
                    instruction = value[m.end():].strip()
                else:
                    instruction = value
                user_text = instruction
            elif c["from"] == "assistant":
                asst_text = c["value"]

        # Build full text sequence
        if self.model_type == "os_atlas":
            full_text = build_internvl_full(user_text, asst_text)
            prompt_text = build_internvl_prompt(user_text)
        else:
            # SeeClick / Qwen-VL format (simplified)
            full_text = f"User: {user_text}\nAssistant: {asst_text}"
            prompt_text = f"User: {user_text}\nAssistant: "

        # Tokenize
        full_ids = self.tokenizer.encode(full_text, add_special_tokens=True)
        prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=True)

        # Truncate
        full_ids = full_ids[: self.max_length]
        n_prompt = min(len(prompt_ids), len(full_ids))

        input_ids = torch.tensor(full_ids, dtype=torch.long)
        labels = torch.tensor(full_ids, dtype=torch.long)
        labels[:n_prompt] = -100  # mask prompt tokens

        # Load image
        pixel_values = self._load_pixel_values(img_path)

        return {
            "input_ids": input_ids,
            "labels": labels,
            "pixel_values": pixel_values,
            "img_path": img_path,
        }

    def _load_pixel_values(self, img_path: str) -> torch.Tensor:
        """Load image and convert to tensor. Returns zeros on failure."""
        try:
            from PIL import Image as PILImage
            import torchvision.transforms as T
            from torchvision.transforms.functional import InterpolationMode

            img = PILImage.open(img_path).convert("RGB")
            transform = T.Compose([
                T.Resize((448, 448), interpolation=InterpolationMode.BICUBIC),
                T.ToTensor(),
                T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            ])
            return transform(img)
        except Exception as exc:
            logger.warning("Failed to load image %s: %s", img_path, exc)
            return torch.zeros(3, 448, 448)