> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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                                                                |
| ------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- |
| [`BagelDatasetConfig`](#nemo_automodel-components-datasets-multimodal-datasets-BagelDatasetConfig)           | Declarative configuration for the packed BAGEL multimodal dataset.         |
| [`SftJSONLIterableDataset`](#nemo_automodel-components-datasets-multimodal-datasets-SftJSONLIterableDataset) | Iterable over conversation JSONL rows where each row may reference images. |
| [`T2IIterableDataset`](#nemo_automodel-components-datasets-multimodal-datasets-T2IIterableDataset)           | Iterable over a parquet-sharded (image, captions) text-to-image dataset.   |

### Functions

| Name                                                                                                                     | Description                                                   |
| ------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------- |
| [`make_bagel_multimodal_dataset`](#nemo_automodel-components-datasets-multimodal-datasets-make_bagel_multimodal_dataset) | Build a BAGEL packed dataset for Stage 1 or Stage 2 training. |

### Data

[`DATASET_REGISTRY`](#nemo_automodel-components-datasets-multimodal-datasets-DATASET_REGISTRY)

[`DEFAULT_DATASET_INFO`](#nemo_automodel-components-datasets-multimodal-datasets-DEFAULT_DATASET_INFO)

[`_MaximumDecompressedSize`](#nemo_automodel-components-datasets-multimodal-datasets-_MaximumDecompressedSize)

[`_MegaByte`](#nemo_automodel-components-datasets-multimodal-datasets-_MegaByte)

[`logger`](#nemo_automodel-components-datasets-multimodal-datasets-logger)

### API

```python
class nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig(
    grouped_datasets: collections.abc.Mapping[str, object] | None = None,
    grouped_datasets_path: str | None = None,
    dataset_info: collections.abc.Mapping[str, object] | None = None,
    dataset_info_path: str | None = None,
    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_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
)
```

Dataclass

Declarative configuration for the packed BAGEL multimodal dataset.

```python
nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig._resolve_mapping(
    inline: collections.abc.Mapping[str, object] | None,
    path: str | None,
    name: str
) -> dict[str, object]
```

staticmethod

Resolve one inline or YAML/JSON mapping without mutating config state.

```python
nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig.build(
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, object],
    rank: int,
    world_size: int,
    num_workers: int,
    global_seed: int,
    data_status: object | None = None
) -> 'PackedDataset'
```

Build and initialize the packed BAGEL dataset.

**Parameters:**

Runtime tokenizer used by the packed dataset.

Runtime BAGEL special-token mapping.

Global distributed rank used for source-data sharding.

Global distributed world size.

Dataloader worker count used to shard source files.

Ranked-training base seed used by worker-start reseeding.

Optional dataset resume state.

**Returns:** `'PackedDataset'`

Initialized packed BAGEL iterable dataset.

```python
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](/nemo-automodel/nemo_automodel/components/datasets/multimodal/distributed_iterable#nemo_automodel-components-datasets-multimodal-distributed_iterable-DistributedIterableDataset)

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

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

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

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

```python
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](/nemo-automodel/nemo_automodel/components/datasets/multimodal/distributed_iterable#nemo_automodel-components-datasets-multimodal-distributed_iterable-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.

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

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

```python
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).

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

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

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

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

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