File size: 3,112 Bytes
fd448dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """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
|