nemo_automodel.components.datasets.multimodal.datasets

BAGEL multimodal datasets: VLM-SFT, T2I pretrain, unified image editing.

The three dataset families emit the data-yield surface needed for the 3-group example.yaml mixture to be iterated by AM’s :class:.packing.PackedDataset. Stage 1 consumes only the understanding-side loss-bearing pieces; Stage 2 additionally consumes VAE-image plan entries for flow-matching loss. VAE encode, MSE computation, noise sampling, and timestep embedding are intentionally training/model concerns, not dataset concerns.

Module Contents

Classes

Name	Description
`SftJSONLIterableDataset`	Iterable over conversation JSONL rows where each row may reference images.
`T2IIterableDataset`	Iterable over a parquet-sharded (image, captions) text-to-image dataset.

Functions

Name	Description
`make_bagel_multimodal_dataset`	Build a BAGEL packed dataset for Stage 1 or Stage 2 training.

Data

DATASET_REGISTRY

DEFAULT_DATASET_INFO

_MaximumDecompressedSize

_MegaByte

logger

API

class nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    frame_sampler,
    jsonl_path_list,
    data_dir_list,
    num_used_data,
    local_rank = 0,
    world_size = 1,
    num_workers = 8,
    data_status = None,
    shuffle_lines = False,
    shuffle_seed = 0
)

Bases: DistributedIterableDataset

Iterable over conversation JSONL rows where each row may reference images.

data_paths

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.change_format(
    data,
    num_images
)

nemo_automodel.components.datasets.multimodal.datasets.SftJSONLIterableDataset.get_data_paths(
    jsonl_path_list,
    data_dir_list,
    num_used_data,
    shuffle_lines,
    shuffle_seed
)

class nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    data_dir_list,
    num_used_data,
    local_rank = 0,
    world_size = 1,
    num_workers = 8,
    data_status = None
)

Bases: DistributedIterableDataset

Iterable over a parquet-sharded (image, captions) text-to-image dataset.

The yielded dict carries only one image tensor (the VAE input). Stage 1 models ignore the VAE branch; Stage 2 consumes it via the PackedDataset vae_image branch + flow-matching loss. No VAE encoding happens in the data pipeline — just tensor preparation.

data_paths

= self.get_data_paths(data_dir_list, num_used_data)

nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset.__iter__()

nemo_automodel.components.datasets.multimodal.datasets.T2IIterableDataset.get_data_paths(
    data_dir_list,
    num_used_data
)

nemo_automodel.components.datasets.multimodal.datasets.make_bagel_multimodal_dataset(
    tokenizer,
    special_tokens,
    grouped_datasets,
    local_rank: int = 0,
    world_size: int = 1,
    num_workers: int = 4,
    expected_num_tokens: int = 32768,
    max_num_tokens_per_sample: int = 16384,
    max_num_tokens: int = 36864,
    prefer_buffer_before: int = 16384,
    max_buffer_size: int = 50,
    interpolate_pos: bool = False,
    use_flex: bool = False,
    data_status = None,
    data_seed: int = 42,
    text_cond_dropout_prob: float = 0.1,
    vit_cond_dropout_prob: float = 0.4,
    vae_cond_dropout_prob: float = 0.1,
    vae_image_downsample: int = 16,
    max_latent_size: int = 32,
    vit_patch_size: int = 14,
    max_num_patch_per_side: int = 70,
    dataset_info = None
)

Build a BAGEL packed dataset for Stage 1 or Stage 2 training.

grouped_datasets is the YAML dict produced by data/configs/*.yaml. dataset_info is required and contains the concrete local paths for each dataset named by grouped_datasets.

Returns the constructed :class:PackedDataset instance. The caller is still responsible for invoking set_epoch immediately before DataLoader iteration (and for the pre-iter RNG reseed — see :meth:PackedDataset.set_epoch docstring).

nemo_automodel.components.datasets.multimodal.datasets.DATASET_REGISTRY = {'t2i_pretrain': T2IIterableDataset, 'vlm_sft': SftJSONLIterableDataset, 'unifie...

nemo_automodel.components.datasets.multimodal.datasets.DEFAULT_DATASET_INFO = {}

nemo_automodel.components.datasets.multimodal.datasets._MaximumDecompressedSize = 1024

nemo_automodel.components.datasets.multimodal.datasets._MegaByte = 2 ** 20

nemo_automodel.components.datasets.multimodal.datasets.logger = logging.getLogger(__name__)