o wiJ@sddlmZmZmZddlmZddlZddl Z ddl m Z m Z ddl mZddlmZmZddlmZddlmZGdd d ejZGd d d eZdS) )DictListOptionalN)EVAL_DATALOADERSTRAIN_DATALOADERS)data) DataLoaderDataset)MegatronDataSampler)loggingcseZdZdZ              d'd ed eed ed ed ededeeedededededededeffdd Zd(de ddfddZ de fddZ de fd d!Zde fd"d#Zd$edefd%d&ZZS))MockDataModuleaA mock LightningDataModule for generating synthetic data for Llama4 models. This module creates dummy datasets (train, validation, test) using `MockLlama4Dataset` for testing or development purposes without requiring actual data. Args: seq_length (int): The sequence length for text tokens. Defaults to 2048. decoder_seq_length (Optional[int]): The sequence length for the decoder (if applicable). Defaults to None. tokenizer (Optional): Tokenizer object. image_processor (Optional): Image processor object. micro_batch_size (int): Micro batch size per GPU. Defaults to 4. global_batch_size (int): Global batch size across all GPUs. Defaults to 8. rampup_batch_size (Optional[List[int]]): Ramp-up schedule for batch size. Defaults to None. num_train_samples (int): Number of synthetic samples for the training set. Defaults to 10,000,000. num_val_samples (int): Number of synthetic samples for the validation set. Defaults to 10,000,000. num_test_samples (int): Number of synthetic samples for the test set. Defaults to 10,000,000. num_workers (int): Number of worker processes for data loading. Defaults to 8. pin_memory (bool): Whether to pin memory for faster data transfer to GPU. Defaults to True. persistent_workers (bool): Whether to keep worker processes alive between epochs. Defaults to False. packed_sequence (bool): Whether to use packed sequences for efficiency. Defaults to False. N逖TF seq_lengthdecoder_seq_length tokenizerimage_processormicro_batch_sizeglobal_batch_sizerampup_batch_sizenum_train_samplesnum_val_samplesnum_test_samples num_workers pin_memorypersistent_workerspacked_sequencecst||_||_||_||_||_| |_| |_| |_ | |_ | |_ ||_ |dus.|durQt dddlm}ddlm}|d}|pI|d|_|pO|j|_t|j|j|||d|_dS)a=A mock LightningDataModule for generating synthetic data for Llama4 models. This module creates dummy datasets (train, validation, test) using `MockLlama4Dataset` for testing or development purposes without requiring actual data. Args: seq_length (int): The sequence length for text tokens. Defaults to 2048. decoder_seq_length (Optional[int]): The sequence length for the decoder (if applicable). Defaults to None. tokenizer (Optional): Tokenizer object. image_processor (Optional): Image processor object. micro_batch_size (int): Micro batch size per GPU. Defaults to 4. global_batch_size (int): Global batch size across all GPUs. Defaults to 8. rampup_batch_size (Optional[List[int]]): Ramp-up schedule for batch size. Defaults to None. num_train_samples (int): Number of synthetic samples for the training set. Defaults to 10,000,000. num_val_samples (int): Number of synthetic samples for the validation set. Defaults to 10,000,000. num_test_samples (int): Number of synthetic samples for the test set. Defaults to 10,000,000. num_workers (int): Number of worker processes for data loading. Defaults to 8. pin_memory (bool): Whether to pin memory for faster data transfer to GPU. Defaults to True. persistent_workers (bool): Whether to keep worker processes alive between epochs. Defaults to False. packed_sequence (bool): Whether to use packed sequences for efficiency. Defaults to False. NzdProcessor or tokenizer are not provided! Fall back to `'meta-llama/Llama-4-Scout-17B-16E-Instruct'`.r) AutoProcessor) AutoTokenizerz)meta-llama/Llama-4-Scout-17B-16E-Instruct)seq_lendecoder_seq_lenrrr)super__init__rr"rrrrrrrrrr warning transformersr=nemo.collections.common.tokenizers.huggingface.auto_tokenizerr from_pretrainedrrr data_sampler)selfrrrrrrrrrrrrrrrr processor __class__b/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/nemo/collections/vlm/llama4/data/mock.pyr$3s8 &     zMockDataModule.__init__stagereturncCs|j}|jr|jdkr||j}td|dt|j|jd|j||jd|_ t|j|jd|j ||jd|_ t|j|jd|j ||jd|_ dS) alSets up the mock datasets for the specified stage. Initializes `MockLlama4Dataset` instances for train, validation, and test splits. Adjusts sequence length if packed sequences are used. Args: stage (str): The stage for which to set up data ('fit', 'validate', 'test', or ''). Defaults to "". zPacked sequence is used with mock dataset. Sequence length for each sample is update to `seq_length // self.micro_batch_size = z`!train)rvalidtestN)rrrr r%MockLlama4Datasetrrr _train_dsr_validation_dsr_test_ds)r*r1rr.r.r/setupxs@   zMockDataModule.setupcC t|ds |d||jS)z,Returns the DataLoader for the training set.r8fit)hasattrr;_create_dataloaderr8r*r.r.r/train_dataloader   zMockDataModule.train_dataloadercCr<)z.Returns the DataLoader for the validation set.r9validate)r>r;r?r9r@r.r.r/val_dataloaderrBzMockDataModule.val_dataloadercCr<)z(Returns the DataLoader for the test set.r:r6)r>r;r?r:r@r.r.r/test_dataloaderrBzMockDataModule.test_dataloaderdatasetcKs"t|f|j|j|j|jd|S)a&Creates a DataLoader for the given dataset. Args: dataset (Dataset): The dataset to wrap in a DataLoader. **kwargs: Additional arguments passed to the DataLoader constructor. Returns: DataLoader: The configured DataLoader instance. )rrr collate_fn)rrrrrG)r*rFkwargsr.r.r/r?s z!MockDataModule._create_dataloader)r NNNrrNrrrrTFF)r0)__name__ __module__ __qualname____doc__intrrboolr$strr;rrArrDrEr rr? __classcell__r.r.r,r/r sd      E*r cseZdZdZ     ddeded ed ed ed ed dffdd Zd efddZ ded e j fddZ ded e eejffddZdee eejfd e eeejffddZddZZS)r7a9A mock Dataset implementation for generating synthetic Llama4 data. Produces batches containing dummy image tensors and random token sequences, mimicking the structure expected by Llama4 models. Args: tokenizer: Tokenizer object to determine vocabulary size. image_processor: Image processor object to determine image dimensions. name (str): Name of the dataset split (e.g., "train", "valid", "test"). num_samples (int): Total number of synthetic samples in this dataset. seq_length (int): Sequence length for the generated token sequences. seed (int): Random seed for data generation reproducibility. Defaults to 42. packed_sequence (bool): Whether the data should be formatted for packed sequences. Defaults to False. pixel_shuffle_ratio (float): Ratio used for calculating the image sequence length after potential pixel shuffling. Defaults to 0.5. num_image_embeddings_per_tile (int): Number of embeddings produced per image tile by the vision encoder (before pixel shuffle). Defaults to 576. num_tiles_per_image (int): Number of tiles the image is split into. Defaults to 1. *F?@r3name num_samplesrseedrpixel_shuffle_ratior2Nc st||_||_||_||_d|_|j} | d| d|_|_ ||_ ||_ ||_ | |_ | |_||_t| | |||_tj|jtjd|_tj|jtjd|_dS)aA mock Dataset implementation for generating synthetic Llama4 data. Produces batches containing dummy image tensors and random token sequences, mimicking the structure expected by Llama4 models. Args: tokenizer: Tokenizer object to determine vocabulary size. image_processor: Image processor object to determine image dimensions. name (str): Name of the dataset split (e.g., "train", "valid", "test"). num_samples (int): Total number of synthetic samples in this dataset. seq_length (int): Sequence length for the generated token sequences. seed (int): Random seed for data generation reproducibility. Defaults to 42. packed_sequence (bool): Whether the data should be formatted for packed sequences. Defaults to False. pixel_shuffle_ratio (float): Ratio used for calculating the image sequence length after potential pixel shuffling. Defaults to 0.5. num_image_embeddings_per_tile (int): Number of embeddings produced per image tile by the vision encoder (before pixel shuffle). Defaults to 576. num_tiles_per_image (int): Number of tiles the image is split into. Defaults to 1. i@ heightwidth)dtypeN)r#r$rTrrr vocab_sizesize image_height image_widthlengthrVrnum_image_embeddings_per_tilenum_tiles_per_imagerWrM _img_seq_lentorchonesfloat loss_maskarangeint64 position_ids) r*rrrTrUrrVrrWr`rar\r,r.r/r$s& "zMockLlama4Dataset.__init__cCs|jS)z-Returns the number of samples in the dataset.)r_r@r.r.r/__len__szMockLlama4Dataset.__len__idxcCs,tjj|j|d}|j|j|jgtjdS)zGenerates a random sequence of token IDs (unused in current __getitem__). Args: idx (int): Index of the sample, used for seeding the random generator. Returns: np.ndarray: An array of random token IDs. rVr\rZ)nprandom default_rngrVintegersr[rrh)r*rknp_genr.r.r/ _get_texts zMockLlama4Dataset._get_textcCstjj|j|d}t|j|j|jdgtj d}d|dd|j <| }t|j|j d|j |jgtjd}|dd}|dd}||||j|jd S) aGenerates a single synthetic data sample. Creates random tensors for 'media' (images), 'tokens', and 'labels'. The 'tokens' sequence includes placeholder IDs where image features would normally be inserted. Args: idx (int): Index of the sample to generate. Used for seeding. Returns: Dict[str, torch.Tensor]: A dictionary containing: - "media": A dummy image tensor [num_tiles, 3, H, W]. - "tokens": Input token sequence [seq_length]. - "labels": Target token sequence (shifted tokens) [seq_length]. - "loss_mask": Mask indicating which tokens contribute to loss [seq_length]. - "position_ids": Positional IDs for the sequence [seq_length]. rlr3rmi N)mediatokenslabelsrfri)rnrorprVrc from_numpyrqr[rrhrbclonerar]r^float32bfloat16rfri)r*rkrrrxryimagesr.r.r/ __getitem__)s4    zMockLlama4Dataset.__getitem__batchc Cstj|}d|d<|djdg|djddR|d<|jrsddlm}|d}|jd}|j}t j d|d ||t j |j d }t j d|d ||t j |j d }d } |||||||| d } | |d <dD] } || d d|| <qf|S)aYCollates a batch of samples from the dataset. Uses the default PyTorch collate function and then performs specific adjustments: - Sets 'attention_mask' to None. - Reshapes 'media' tensor. - If `packed_sequence` is True, it prepares the batch for packed sequence format by calculating cumulative sequence lengths (`cu_seqlens`) and reshaping relevant tensors ('tokens', 'labels', 'loss_mask', 'position_ids'). Args: batch (List[Dict[str, torch.Tensor]]): A list of individual samples (dictionaries) from `__getitem__`. Returns: Dict[str, Optional[torch.Tensor]]: The collated batch, ready for model input. Includes 'packed_seq_params' if packing is enabled. Nattention_maskrwrvrtr)PackedSeqParamsrxr3)steprZdevicethd) cu_seqlens_q cu_seqlens_kvcu_seqlens_q_paddedcu_seqlens_kv_padded max_seqlen_q max_seqlen_kv qkv_formatpacked_seq_params)rxryrfri) r dataloaderdefault_collatereshapeshapermegatron.core.packed_seq_paramsrrrcrgint32r) r*rcollated_batchrrx batch_size valid_seqlen cu_seqlenscu_seqlens_paddedrrkeyr.r.r/ _collate_fn[s8 (   zMockLlama4Dataset._collate_fncCs ||S)aMethod passed to the DataLoader's `collate_fn` argument. Simply calls the internal `_collate_fn` implementation. This structure allows for potential future additions like neural type checking within this wrapper method. Args: batch: A list of samples fetched from the dataset. Returns: The collated batch dictionary. )r)r*rr.r.r/rGs zMockLlama4Dataset.collate_fn)rQFrRrSr3)rIrJrKrLrOrMrNrer$rjrnndarrayrsrrcTensorrrrrrGrPr.r.r,r/r7s6  : .24r7)typingrrrlightning.pytorchpytorchplnumpyrnrc!lightning.pytorch.utilities.typesrr torch.utilsrtorch.utils.datarr nemo.lightning.pytorch.pluginsr nemo.utilsr LightningDataModuler r7r.r.r.r/s    -