"""Utilities for multimodal VLM datasets.""" import json import os from pathlib import Path from typing import Any, Dict, List, Tuple from PIL import Image from torchvision import transforms from transformers import AutoTokenizer from taoTrain.config import TrainingConfig from taoTrain.data.sft_utils import parse_sft_record from taoTrain.data.tokenizer import ( SentencePieceTokenizerWrapper, load_special_token_metadata, require_special_token_id, ) def resolve_special_token_id( tokenizer, token_name: str, token_value: str, tokenizer_path: str | None = None, ) -> int: """Resolve a configured special token and enforce the single-token contract.""" try: return require_special_token_id(tokenizer, token_value) except ValueError as exc: tokenizer_hint = f" in tokenizer `{tokenizer_path}`" if tokenizer_path else "" raise ValueError(f"{token_name} `{token_value}` could not be resolved{tokenizer_hint}: {exc}") from exc def validate_vlm_special_tokens(config: TrainingConfig, tokenizer) -> Dict[str, int]: """Validate and resolve the reserved special tokens required by multimodal training.""" tokenizer_path = config.dataset.tokenizer_path return { "image_token_id": resolve_special_token_id(tokenizer, "image_token", config.image_token, tokenizer_path), "user_token_id": resolve_special_token_id(tokenizer, "user_token", config.user_token, tokenizer_path), "assistant_token_id": resolve_special_token_id( tokenizer, "assistant_token", config.assistant_token, tokenizer_path, ), } def load_tokenizer_from_config(config: TrainingConfig): """Load a tokenizer using the same resolution rules as the JSONL datasets.""" dataset_config = config.dataset if dataset_config.tokenizer_path: tokenizer_path = Path(dataset_config.tokenizer_path) tokenizer_type = dataset_config.tokenizer_type if tokenizer_type is None: tokenizer_type = "sentencepiece" if dataset_config.tokenizer_path.endswith(".model") else "huggingface" if tokenizer_type == "sentencepiece": if not tokenizer_path.exists(): raise FileNotFoundError( f"SentencePiece tokenizer not found: {dataset_config.tokenizer_path}. " "Train or regenerate the tokenizer artifact before running VLM training." ) import sentencepiece as spm processor = spm.SentencePieceProcessor() processor.Load(dataset_config.tokenizer_path) special_token_ids = load_special_token_metadata(dataset_config.tokenizer_path) tokenizer = SentencePieceTokenizerWrapper(processor, special_token_ids=special_token_ids) else: tokenizer = AutoTokenizer.from_pretrained(dataset_config.tokenizer_path) else: tokenizer = AutoTokenizer.from_pretrained("gpt2") if hasattr(tokenizer, "pad_token") and tokenizer.pad_token is None and hasattr(tokenizer, "eos_token"): tokenizer.pad_token = tokenizer.eos_token return tokenizer def load_multimodal_records(config: TrainingConfig) -> list[dict[str, Any]]: """Load local multimodal JSONL records.""" jsonl_path = Path(config.dataset.jsonl_path or "") if not jsonl_path.exists(): raise FileNotFoundError(f"VLM dataset not found: {jsonl_path}") records: list[dict[str, Any]] = [] with jsonl_path.open("r", encoding="utf-8") as handle: for line in handle: line = line.strip() if not line: continue records.append(json.loads(line)) if config.dataset.max_samples and len(records) >= config.dataset.max_samples: break return records def _normalize_image_path_value(image_value: Any, column_name: str) -> str: """Normalize a dataset image field into a single image path string.""" if isinstance(image_value, os.PathLike): return os.fspath(image_value) if isinstance(image_value, str): if image_value.strip(): return image_value raise ValueError(f"Image path field `{column_name}` cannot be an empty string") if isinstance(image_value, (list, tuple)): if not image_value: raise ValueError(f"Image path field `{column_name}` cannot be an empty list") for entry in image_value: if isinstance(entry, os.PathLike): return os.fspath(entry) if isinstance(entry, str) and entry.strip(): return entry if isinstance(entry, (list, tuple)): raise ValueError( f"Image path field `{column_name}` must be a flat list of paths, got nested value: {entry!r}" ) if entry is not None and entry != "": raise ValueError( f"Image path field `{column_name}` must contain only string paths, got {type(entry).__name__}" ) raise ValueError(f"Image path field `{column_name}` does not contain a usable image path") raise ValueError( f"Image path field `{column_name}` must be a string path or list of string paths, " f"got {type(image_value).__name__}" ) def _resolve_image_path(record: Dict[str, Any], config: TrainingConfig) -> str: """Resolve the image path from the configured field or common aliases.""" dataset_config = config.dataset configured_column = dataset_config.image_path_column candidate_columns = [configured_column, *dataset_config.image_path_aliases] field_errors: list[str] = [] for column_name in dict.fromkeys(candidate_columns): if column_name not in record: continue image_path = record.get(column_name) if image_path is None: continue try: return _normalize_image_path_value(image_path, column_name) except ValueError as exc: field_errors.append(str(exc)) available_keys = ", ".join(sorted(record.keys())) or "" error_hint = f" Field errors: {'; '.join(field_errors)}." if field_errors else "" raise ValueError( "Multimodal record must contain an image path field. " f"Configured field: `{configured_column}`. " f"Tried aliases: {dataset_config.image_path_aliases}. " f"Available keys: {available_keys}.{error_hint}" ) def parse_vlm_record(record: Dict[str, Any], config: TrainingConfig) -> Tuple[str, List[Tuple[str, str]]]: """Parse a multimodal record into image path and SFT-like turns.""" image_path = _resolve_image_path(record, config) turns = [] if "turns" in record: turns, _ = parse_sft_record({"turns": record["turns"]}, config) elif "input" in record and "output" in record: turns = [(record["input"], record["output"])] else: dataset_config = config.dataset instruction_col = dataset_config.instruction_column or "instruction" response_col = dataset_config.response_column or "response" if instruction_col in record and response_col in record: turns = [(record[instruction_col], record[response_col])] elif "text" in record and str(record["text"]).strip(): turns = [(config.dataset.caption_prompt, record["text"])] if not turns: available_keys = ", ".join(sorted(record.keys())) or "" raise ValueError( "Multimodal record must contain either `input`/`output`, `turns`, or caption-style `text`. " f"Available keys: {available_keys}" ) return image_path, turns def build_vlm_sequence_tokens( turns: List[Tuple[str, str]], tokenizer, image_token: str, user_token: str, assistant_token: str, max_seq_length: int, ) -> Tuple[List[int], List[int], List[int], int]: """Build a multimodal SFT sequence that starts with a single image placeholder token.""" image_token_id = resolve_special_token_id(tokenizer, "image_token", image_token) user_token_id = resolve_special_token_id(tokenizer, "user_token", user_token) assistant_token_id = resolve_special_token_id(tokenizer, "assistant_token", assistant_token) input_ids = [image_token_id] mask = [0] eos_token_id = require_special_token_id(tokenizer, "") pad_token_id = require_special_token_id(tokenizer, "") for user_text, assistant_text in turns: input_ids.append(user_token_id) mask.append(0) user_tokens = tokenizer(user_text, add_special_tokens=False)["input_ids"] input_ids.extend(user_tokens) mask.extend([0] * len(user_tokens)) input_ids.append(assistant_token_id) mask.append(0) assistant_tokens = tokenizer(assistant_text, add_special_tokens=False)["input_ids"] input_ids.extend(assistant_tokens) mask.extend([1] * len(assistant_tokens)) real_length = len(input_ids) input_ids.append(eos_token_id) mask.append(0) real_length += 1 if len(input_ids) > max_seq_length: input_ids = input_ids[:max_seq_length] mask = mask[:max_seq_length] real_length = max_seq_length padding_len = max_seq_length - len(input_ids) if padding_len > 0: input_ids.extend([pad_token_id] * padding_len) mask.extend([0] * padding_len) attention_mask = [1] * real_length + [0] * max(0, max_seq_length - real_length) return input_ids, attention_mask, mask, image_token_id def build_image_transform(image_size: int): """Build the standard VLM image preprocessing pipeline.""" return transforms.Compose([ transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) def load_image(image_path: str, dataset_path: str, transform) -> Any: """Load an image from a local path, resolving relative paths against the JSONL file.""" image_file = Path(image_path) if not image_file.is_absolute(): image_file = Path(dataset_path).parent / image_file if not image_file.exists(): raise FileNotFoundError(f"Image file not found: {image_file}") image = Image.open(image_file).convert("RGB") return transform(image)