import os import math import random import logging import yaml from typing import Dict, Any, List, Optional from collections import defaultdict import torch from torch.utils.data import Dataset, ConcatDataset from smolvlm.datasets.dataset import SupervisedDataset from smolvlm.train.args import DataArguments, TrainingArguments, ModelArguments from smolvlm.constants import IGNORE_INDEX from smolvlm.utils import mprint # or your custom printing utility from tabulate import tabulate from transformers import ProcessorMixin from torch.nn.utils.rnn import pad_sequence logger = logging.getLogger(__name__) ############################################################################## # Multi-subsequence "varlen" packing with integer-coded `subseq_ids` ############################################################################## def len2weight(num_effective_tokens: int, mode: str) -> float: """ Returns the sub-sample weight given the sub-sequence length. """ if num_effective_tokens == 0: return 0.0 # or skip if mode == "token": return 1.0 # no length-based weighting elif mode == "sample": # each sub-sample counts equally, so 1 / length return 1.0 / num_effective_tokens elif mode == "square": # default in InternVL return 1.0 / (num_effective_tokens**0.5) else: # 'none' or fallback return 1.0 class PackedConcatDataset(ConcatDataset): """ Merges multiple short sub-samples from an underlying ConcatDataset into a single “packed” sample, up to a `cutoff_len` tokens. Assigns integer-coded sub-sequence IDs in `subseq_ids`: 1 => sub-sample #1 2 => sub-sample #2 ... so your collator can turn them into block diagonal (varlen) attention masks. Each returned item from __getitem__ is: { "input_ids": (sum_of_sub_len,) int, "labels": (sum_of_sub_len,) int, "subseq_ids": (sum_of_sub_len,) int in [1..N], "pixel_values": (sum_of_frames, 3, H, W) if images exist } We do NOT do final token-level padding to cutoff_len; we let the data collator handle batch-level padding (less wasteful). Attributes: datasets: List of sub-datasets we are merging. cutoff_len: Max tokens we want in a single “packed” sample. pad_token_id: If needed for partial fix-ups. packed_cursor: Tracks how many sub-samples we have consumed so far. """ def __init__( self, datasets: List, data_args, model_max_length: int = 2048 ): super().__init__(datasets) self.data_args = data_args self.cutoff_len = max(model_max_length, 1) self.pad_token_id = getattr(data_args, "pad_token_id", 0) self.packed_cursor = 0 logger.info( f"[PackedConcatDataset] Using cutoff_len={self.cutoff_len}; " f"we'll merge multiple sub-samples per item." ) def __len__(self): # Underlying total number of sub-samples return super().__len__() def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: """ We ignore 'idx' because each call to __getitem__ gets the “next” packed sample, from self.packed_cursor onward. """ if self.packed_cursor >= len(self): raise IndexError("No more sub-samples left to pack in PackedConcatDataset.") # Accumulate sub-samples chunk_input_ids = [] chunk_labels = [] chunk_subseq_ids = [] pixel_key = None pixel_values_list = [] sub_seq_counter = 0 current_token_count = 0 while True: if self.packed_cursor >= len(self): break sub_item = super().__getitem__(self.packed_cursor) self.packed_cursor += 1 sub_len = sub_item["input_ids"].size(0) if (current_token_count > 0) and (current_token_count + sub_len) > self.cutoff_len: # Revert if we can't fit this sub-sample self.packed_cursor -= 1 break sub_seq_counter += 1 seq_id_tensor = torch.full( (sub_len,), fill_value=sub_seq_counter, dtype=torch.long, device=sub_item["input_ids"].device ) chunk_input_ids.append(sub_item["input_ids"]) chunk_labels.append(sub_item["labels"]) chunk_subseq_ids.append(seq_id_tensor) # If images are present if "pixel_values" in sub_item: pixel_key = "pixel_values" pixel_values_list.append(sub_item["pixel_values"]) current_token_count += sub_len print("[Sequence Packing] current num tokens:", current_token_count) if current_token_count >= self.cutoff_len: break # Merge text if len(chunk_input_ids) == 0: return { "input_ids": torch.tensor([], dtype=torch.long), "labels": torch.tensor([], dtype=torch.long), "attention_mask": torch.tensor([], dtype=torch.long), } merged_input_ids = torch.cat(chunk_input_ids, dim=0) merged_labels = torch.cat(chunk_labels, dim=0) merged_subseq_ids = torch.cat(chunk_subseq_ids, dim=0) # Merge images along frame dimension if present merged_pixel_values = None if pixel_key and pixel_values_list: merged_pixel_values = torch.cat(pixel_values_list, dim=0) # shape => (f1+f2+..., 3, H, W) loss_weight = torch.ones_like(merged_subseq_ids, dtype=torch.float32) unique_ids = merged_subseq_ids.unique() unique_ids = unique_ids[unique_ids > 0] # ignore pad=0 for sid in unique_ids.tolist(): mask = (merged_subseq_ids == sid) num_eff = (merged_labels[mask] != IGNORE_INDEX).sum().item() w = len2weight(num_eff, self.data_args.loss_reduction) loss_weight[mask] = w # Build final out_dict = { "input_ids": merged_input_ids, "labels": merged_labels, "attention_mask": merged_subseq_ids, "loss_weight": loss_weight, } if merged_pixel_values is not None: out_dict[pixel_key] = merged_pixel_values return out_dict # Varlen Collator (subseq_ids => block diagonal) ############################################################################## def pad_sequence_varlen( sequences: List[torch.Tensor], batch_first: bool = True, padding_value: float = 0.0 ) -> torch.Tensor: """ Similar signature to torch.nn.utils.rnn.pad_sequence, but treats each Tensor as a 1D sequence of any integer-coded or float tokens, and uses integer-coded varlen semantics if desired. If batch_first=True, returns (batch_size, seq_len). Otherwise returns (seq_len, batch_size). """ if len(sequences) == 0: # Return an empty tensor if no sequences return torch.tensor([], dtype=torch.long) max_len = max(seq.size(0) for seq in sequences) batch_size = len(sequences) device = sequences[0].device dtype = sequences[0].dtype if batch_first: # Shape => (batch_size, max_len) out = torch.full((batch_size, max_len), padding_value, device=device, dtype=dtype) for i, seq in enumerate(sequences): length = seq.size(0) out[i, :length] = seq else: # Shape => (max_len, batch_size) out = torch.full((max_len, batch_size), padding_value, device=device, dtype=dtype) for i, seq in enumerate(sequences): length = seq.size(0) out[:length, i] = seq return out ############################################################################## # Standard Collator (0/1 attn mask) ############################################################################## class DataCollatorForSupervisedDataset: """ Collates examples containing text-only or text+image/video data. 1) Text sequences (input_ids, attention_mask, labels) are padded to the maximum batch length. - If model_max_length is set, we optionally truncate. 2) Pixel data (pixel_values, optional) is padded to (max_frames, 3, max_h, max_w). 3) Pixel-level mask (pixel_attention_mask): - If provided in the example, we pad it accordingly. - If not provided, we fill the valid image region with 1, the remainder with 0. """ def __init__( self, pad_token_id: int, model_max_length: Optional[int] = None, image_size: Optional[int] = None, func_pad_sequence = pad_sequence, ): self.pad_token_id = pad_token_id self.ignore_index = IGNORE_INDEX self.model_max_length = model_max_length self.image_size = image_size self.func_pad_sequence = func_pad_sequence def __call__(self, examples: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: ################################################################ # PART A: Pad the text data (input_ids, attention_mask, labels) ################################################################ attention_masks_list = [] for ex in examples: # If "attention_mask" is missing, we generate it on the fly if "attention_mask" in ex: attention_masks_list.append(ex["attention_mask"]) else: am = (ex["input_ids"] != self.pad_token_id).long() attention_masks_list.append(am) input_ids = self.func_pad_sequence( [ex["input_ids"] for ex in examples], batch_first=True, padding_value=self.pad_token_id ) attention_mask = self.func_pad_sequence( attention_masks_list, batch_first=True, padding_value=0 ) labels = self.func_pad_sequence( [ex["labels"] for ex in examples], batch_first=True, padding_value=self.ignore_index ) # Optional: truncate if model_max_length is specified if self.model_max_length and input_ids.size(1) > self.model_max_length: input_ids = input_ids[:, :self.model_max_length] attention_mask = attention_mask[:, :self.model_max_length] labels = labels[:, :self.model_max_length] out = { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels } ################################################################ # PART B: Handle pixel data (pixel_values) + pixel_attention_mask ################################################################ # Step 1: figure out maximum frames, height, width across the batch pvs = [ex["pixel_values"] for ex in examples if "pixel_values" in ex] if pvs: # there is at least one non-None pixel_values max_frames = max(pv.shape[0] for pv in pvs) max_h = max(pv.shape[-2] for pv in pvs) max_w = max(pv.shape[-1] for pv in pvs) else: max_h = max_w = self.image_size max_frames = 1 #TODO: verify this is good default # Step 2: create padded pixel_values and pixel_attention_mask for each example padded_pixel_values_list = [] padded_pixel_mask_list = [] for ex in examples: pv = ex.get("pixel_values", None) pm = ex.get("pixel_attention_mask", None) # shape (f, h, w) if provided if pv is None: # text-only => fill pixel data + mask with zeros shape_pv = (max_frames, 3, max_h, max_w) shape_pm = (max_frames, max_h, max_w) padded_pv = torch.zeros(shape_pv, dtype=torch.float32) padded_pm = torch.zeros(shape_pm, dtype=torch.long) else: f, c, h, w = pv.shape # Prepare final storage padded_pv = torch.zeros( (max_frames, c, max_h, max_w), dtype=pv.dtype, device=pv.device ) padded_pm = torch.zeros( (max_frames, max_h, max_w), dtype=torch.long, device=pv.device ) padded_pv[:f, :, :h, :w] = pv # Copy or fill the pixel attention mask if pm is not None: padded_pm[:f, :h, :w] = pm else: # Mark valid region as 1 padded_pm[:f, :h, :w] = 1 padded_pixel_values_list.append(padded_pv) padded_pixel_mask_list.append(padded_pm) # Finally, stack along batch dimension ## try not outputting pixel_values in text-only sample #if any("pixel_values" in ex for ex in examples): out["pixel_values"] = torch.stack(padded_pixel_values_list, dim=0) return out ############################################################################## # Summaries ############################################################################## def display_overview(summary_data: Dict[str, Any], total_count: int) -> None: print("=== Overview ===") print(f"Aggregate Sample Count: {total_count}\n") for category, info in summary_data.items(): ctotal = info["total_samples"] cpct = (ctotal / total_count * 100) if total_count > 0 else 0 print(f"{category.title()} Overview") print(f"Number of Samples: {ctotal} ({cpct:.2f}%)") print("-" * 50) table_data = [] headers = ["Dataset", "Count", "Percentage"] for entry in info["datasets"]: esamples = entry["samples"] epct = (esamples / total_count * 100) if total_count > 0 else 0 table_data.append([entry["dataset_name"], esamples, f"{epct:.2f}%"]) print(tabulate(table_data, headers=headers, tablefmt="fancy_grid")) print() ############################################################################## # Main builder logic ############################################################################## def build_datasets( data_args: DataArguments, training_args: TrainingArguments, model_args: ModelArguments, processor: ProcessorMixin, split: str = "train", ): """ 1) Load a YAML describing multiple sub-datasets. 2) Create a list of SupervisedDataset objects. 3) If data_args.packed => use PackedConcatDataset (with subseq_ids), else => normal ConcatDataset. """ if getattr(model_args, "frames_per_clip", 1) > 1: from smolvlm.datasets.dataset_clip_sampling import SupervisedDataset else: from smolvlm.datasets.dataset import SupervisedDataset mprint(f"[Dataset-INFO]: Loading from {data_args.data_mixture}") with open(data_args.data_mixture, "r") as yf: meta_datasets = yaml.safe_load(yf) all_datasets = [] extra_info = [] for dataset_type, dataset_list in meta_datasets.items(): for ds_args in dataset_list: ds = SupervisedDataset( dataset_args=ds_args, processor=processor, data_args=data_args, training_args=training_args, model_args=model_args, ) all_datasets.append(ds) extra_info.append({ "dataset_name": ds.name, "modality": ds.modality, "samples": len(ds), }) # Summaries from collections import defaultdict modality_summary = defaultdict(lambda: {"total_samples": 0, "datasets": []}) total_samples = 0 for entry in extra_info: mod = entry["modality"] dsname = entry["dataset_name"] samples = entry["samples"] modality_summary[mod]["total_samples"] += samples modality_summary[mod]["datasets"].append(entry) total_samples += samples display_overview(modality_summary, total_samples) # Build final dataset if data_args.packed: mprint("[build_datasets] Using PackedConcatDataset for multi-sample packing with subseq_ids.") dataset = PackedConcatDataset( all_datasets, data_args=data_args, model_max_length=training_args.model_max_length ) else: mprint("[build_datasets] Using standard ConcatDataset (no packing).") dataset = ConcatDataset(all_datasets) # Save some info in training_args training_args.sample_lens = [e["samples"] for e in extra_info] training_args.data_info = extra_info return dataset def make_supervised_data_module( processor: ProcessorMixin, data_args: DataArguments, training_args: TrainingArguments, model_args: ModelArguments, ): """ Creates train_dataset, eval_dataset, and data_collator. If data_args.packed => we do integer-coded subseq_ids approach, else => normal approach with 0/1 attention_mask. """ train_dataset = build_datasets(data_args, training_args, model_args, processor, split="train") eval_dataset = None if data_args.packed: # Use the varlen collator data_collator = DataCollatorForSupervisedDataset( pad_token_id = processor.tokenizer.pad_token_id, model_max_length = processor.tokenizer.model_max_length, image_size = data_args.video_target_size, func_pad_sequence = pad_sequence_varlen, ) else: # Use the normal collator data_collator = DataCollatorForSupervisedDataset( pad_token_id = processor.tokenizer.pad_token_id, model_max_length = processor.tokenizer.model_max_length, image_size = data_args.video_target_size, func_pad_sequence = pad_sequence, ) return dict( train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=data_collator, )