| """Local JSONL datasets for multimodal VLM training."""
|
|
|
| from typing import Dict
|
| import torch
|
| from torch.utils.data import Dataset
|
|
|
| from taoTrain.data.vlm_utils import ( |
| build_image_transform, |
| build_vlm_sequence_tokens, |
| load_image, |
| load_multimodal_records, |
| load_tokenizer_from_config, |
| parse_vlm_record, |
| validate_vlm_special_tokens, |
| ) |
|
|
|
|
| class VLMJSONLDataset(Dataset):
|
| """JSONL dataset for multimodal connector training and multimodal SFT."""
|
|
|
| def __init__(self, config, split: str = "train"):
|
| """Initialize the multimodal dataset."""
|
| self.config = config
|
| self.split = split
|
| self.records = load_multimodal_records(config) |
| self.tokenizer = load_tokenizer_from_config(config) |
| self.special_token_ids = validate_vlm_special_tokens(config, self.tokenizer) |
| self.transform = build_image_transform(config.image_size) |
| self.text_seq_length = config.model.max_seq_length - config.vision_prefix_tokens + 1 |
|
|
| if self.text_seq_length < 2:
|
| raise ValueError(
|
| "model.max_seq_length must be at least vision_prefix_tokens + 1 for multimodal expansion"
|
| )
|
|
|
| def __len__(self) -> int:
|
| """Return dataset size."""
|
| return len(self.records)
|
|
|
| def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
|
| """Load one multimodal training example."""
|
| record = self.records[idx]
|
| image_path, turns = parse_vlm_record(record, self.config)
|
| pixel_values = load_image(image_path, self.config.dataset.jsonl_path, self.transform)
|
|
|
| input_ids, attention_mask, mask, image_token_id = build_vlm_sequence_tokens(
|
| turns=turns,
|
| tokenizer=self.tokenizer,
|
| image_token=self.config.image_token,
|
| user_token=self.config.user_token,
|
| assistant_token=self.config.assistant_token,
|
| max_seq_length=self.text_seq_length,
|
| )
|
|
|
| first_non_pad_idx = next((i for i, value in enumerate(attention_mask) if value == 1), None)
|
| if first_non_pad_idx != 0 or input_ids[0] != image_token_id:
|
| raise ValueError("Multimodal samples must begin with the configured <image> token")
|
|
|
| labels = input_ids[1:].copy() + [-100] |
| for token_idx, mask_value in enumerate(mask): |
| if mask_value == 0: |
| labels[token_idx] = -100 |
|
|
| if all(label == -100 for label in labels): |
| raise ValueError( |
| f"VLM sample at index {idx} produced no trainable assistant/caption tokens. " |
| "Check the record format and special-token masking." |
| ) |
|
|
| return { |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "attention_mask": torch.tensor(attention_mask, dtype=torch.long), |
| "labels": torch.tensor(labels, dtype=torch.long), |
| "pixel_values": pixel_values,
|
| }
|
|
|
|
|
| class VLMSFTJSONLDataset(VLMJSONLDataset):
|
| """JSONL dataset for end-to-end multimodal supervised fine-tuning."""
|
|
|
| pass
|
|
|
