""" DataCollator for OS-Atlas (InternVL2) training. Handles batching of image + text pairs using InternVL2's conversation format: <|im_start|>user 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 = "" 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"(.*?)", 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)