> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

BAGEL-style multimodal data pipeline for packed three-group training.

Ports the subset of upstream BAGEL `data` needed to feed BAGEL training from
fully AM-native code: VLM SFT, T2I pretrain, and unified image editing. The
packed batch schema is shared by Stage 1 and Stage 2; whether VAE /
flow-matching tensors are consumed is controlled by the model stage.

## Submodules

* **[`nemo_automodel.components.datasets.multimodal.collate_fns`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/collate_fns)**
* **[`nemo_automodel.components.datasets.multimodal.datasets`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/datasets)**
* **[`nemo_automodel.components.datasets.multimodal.distributed_iterable`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/distributed_iterable)**
* **[`nemo_automodel.components.datasets.multimodal.interleave`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/interleave)**
* **[`nemo_automodel.components.datasets.multimodal.loader`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/loader)**
* **[`nemo_automodel.components.datasets.multimodal.packing`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/packing)**
* **[`nemo_automodel.components.datasets.multimodal.parquet_utils`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/parquet_utils)**
* **[`nemo_automodel.components.datasets.multimodal.transforms`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/transforms)**
* **[`nemo_automodel.components.datasets.multimodal.utils`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/utils)**
* **[`nemo_automodel.components.datasets.multimodal.video`](/nemo-automodel/nemo_automodel/components/datasets/multimodal/video)**

## Package Contents

### Classes

| Name                                                                                                                           | Description                                                                              |
| ------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------- |
| [`BagelDatasetConfig`](#nemo_automodel-components-datasets-multimodal-datasets-BagelDatasetConfig)                             | Declarative configuration for the packed BAGEL multimodal dataset.                       |
| [`DistributedIterableDataset`](#nemo_automodel-components-datasets-multimodal-distributed_iterable-DistributedIterableDataset) | Base class for rank/worker-aware iterable datasets.                                      |
| [`FrameSampler`](#nemo_automodel-components-datasets-multimodal-video-FrameSampler)                                            | Callable that returns a list of PIL frames for a given file path.                        |
| [`ImageTransform`](#nemo_automodel-components-datasets-multimodal-transforms-ImageTransform)                                   | Full BAGEL image transform: resize + to\_tensor + normalize.                             |
| [`InterleavedBaseIterableDataset`](#nemo_automodel-components-datasets-multimodal-interleave-InterleavedBaseIterableDataset)   | Builder mixin for interleaved image/text/video sequence plans.                           |
| [`MaxLongEdgeMinShortEdgeResize`](#nemo_automodel-components-datasets-multimodal-transforms-MaxLongEdgeMinShortEdgeResize)     | Resize so longest/shortest edges stay within bounds and both edges are stride-divisible. |
| [`PackedDataset`](#nemo_automodel-components-datasets-multimodal-packing-PackedDataset)                                        | Greedy pack of samples drawn from weighted groups into token-budgeted batches.           |
| [`ParquetStandardIterableDataset`](#nemo_automodel-components-datasets-multimodal-interleave-ParquetStandardIterableDataset)   | Base class: iterate per-(file, row\_group) across a list of parquet shards.              |
| [`SftJSONLIterableDataset`](#nemo_automodel-components-datasets-multimodal-datasets-SftJSONLIterableDataset)                   | Iterable over conversation JSONL rows where each row may reference images.               |
| [`SimpleCustomBatch`](#nemo_automodel-components-datasets-multimodal-collate_fns-SimpleCustomBatch)                            | Pass-through wrapper around one packed batch from `PackedDataset`.                       |
| [`T2IIterableDataset`](#nemo_automodel-components-datasets-multimodal-datasets-T2IIterableDataset)                             | Iterable over a parquet-sharded (image, captions) text-to-image dataset.                 |
| [`UnifiedEditIterableDataset`](#nemo_automodel-components-datasets-multimodal-interleave-UnifiedEditIterableDataset)           | Image-editing dataset: `(input, instruction, output)` chains over parquet.               |

### Functions

| Name                                                                                                                                    | Description                                                                   |
| --------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`add_special_tokens`](#nemo_automodel-components-datasets-multimodal-utils-add_special_tokens)                                         | Add BAGEL's four special tokens to a Qwen2 tokenizer and return the ids.      |
| [`bagel_packed_collate_fn`](#nemo_automodel-components-datasets-multimodal-collate_fns-bagel_packed_collate_fn)                         | Canonical name in AM's collate-fn registry.                                   |
| [`collate_wrapper`](#nemo_automodel-components-datasets-multimodal-collate_fns-collate_wrapper)                                         | Return the BAGEL-style identity collate (wraps a single packed dict).         |
| [`get_flattened_position_ids_extrapolate`](#nemo_automodel-components-datasets-multimodal-utils-get_flattened_position_ids_extrapolate) | Return flattened 2D patch position IDs by direct grid extrapolation.          |
| [`get_flattened_position_ids_interpolate`](#nemo_automodel-components-datasets-multimodal-utils-get_flattened_position_ids_interpolate) | Return flattened 2D patch position IDs by interpolating to the max grid.      |
| [`get_parquet_data_paths`](#nemo_automodel-components-datasets-multimodal-parquet_utils-get_parquet_data_paths)                         | Return a flat list of parquet file paths sharded across ranks.                |
| [`init_arrow_pf_fs`](#nemo_automodel-components-datasets-multimodal-parquet_utils-init_arrow_pf_fs)                                     | Return a pyarrow filesystem for `parquet_file_path`.                          |
| [`len2weight`](#nemo_automodel-components-datasets-multimodal-utils-len2weight)                                                         | Convert a sequence length into BAGEL's per-sample loss weight.                |
| [`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.                 |
| [`patchify`](#nemo_automodel-components-datasets-multimodal-utils-patchify)                                                             | Patchify a CxHxW image tensor into (H/p \* W/p, p*p*C).                       |
| [`pil_img2rgb`](#nemo_automodel-components-datasets-multimodal-utils-pil_img2rgb)                                                       | Convert a PIL image to RGB, compositing transparent pixels on white.          |
| [`prepare_attention_mask_per_sample`](#nemo_automodel-components-datasets-multimodal-utils-prepare_attention_mask_per_sample)           | Build a per-sample additive float mask honoring causal / full / noise splits. |

### Data

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

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

[`DataConfig`](#nemo_automodel-components-datasets-multimodal-packing-DataConfig)

### 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.

**`data_seed`** `int = 42`

---

**`dataset_info`** `Mapping[str, object] | None = None`

---

**`dataset_info_path`** `str | None = None`

---

**`expected_num_tokens`** `int = 32768`

---

**`grouped_datasets`** `Mapping[str, object] | None = None`

---

**`grouped_datasets_path`** `str | None = None`

---

**`interpolate_pos`** `bool = False`

---

**`max_buffer_size`** `int = 50`

---

**`max_latent_size`** `int = 32`

---

**`max_num_patch_per_side`** `int = 70`

---

**`max_num_tokens`** `int = 36864`

---

**`max_num_tokens_per_sample`** `int = 16384`

---

**`prefer_buffer_before`** `int = 16384`

---

**`text_cond_dropout_prob`** `float = 0.1`

---

**`use_flex`** `bool = False`

---

**`vae_cond_dropout_prob`** `float = 0.1`

---

**`vae_image_downsample`** `int = 16`

---

**`vit_cond_dropout_prob`** `float = 0.4`

---

**`vit_patch_size`** `int = 14`

---

```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:**

**`tokenizer`** `'PreTrainedTokenizerBase'`

Runtime tokenizer used by the packed dataset.

---

**`special_tokens`** `Mapping[str, object]`

Runtime BAGEL special-token mapping.

---

**`rank`** `int`

Global distributed rank used for source-data sharding.

---

**`world_size`** `int`

Global distributed world size.

---

**`num_workers`** `int`

Dataloader worker count used to shard source files.

---

**`global_seed`** `int`

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

---

**`data_status`** `object | None` — default: None

Optional dataset resume state.

---

**Returns:** `'PackedDataset'`

Initialized packed BAGEL iterable dataset.

```python
class nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset(
    dataset_name,
    local_rank = 0,
    world_size = 1,
    num_workers = 8
)
```

**Bases:** `IterableDataset`

Base class for rank/worker-aware iterable datasets.

Owns a private `rng` used only to shuffle file paths deterministically
in `set_epoch` — NOT used for per-sample randomness. Per-sample
randomness still goes through the Python global `random` module (see
`packing` for the reseed hook).

**`_drop_counters`** `= {}`

---

**`data_paths_per_rank`** `= []`

---

**`num_files_per_rank`** `= 0`

---

**`rng`** `= random.Random()`

---

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.__iter__()
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._get_worker_data_status(
    worker_id
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._log_drop(
    reason,
    message,
    args = (),
    every = 100,
    exc_info = False
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset._set_worker_resume_data_status(
    worker_id,
    status
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.get_data_paths(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.get_data_paths_per_worker()
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.load_state_dict(
    state_dict
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.set_data_status(
    data_status
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.set_epoch(
    seed = 42
)
```

```python
nemo_automodel.components.datasets.multimodal.distributed_iterable.DistributedIterableDataset.state_dict()
```

```python
class nemo_automodel.components.datasets.multimodal.video.FrameSampler(
    max_num_frames = -1,
    min_num_frames = 8,
    sample = 'rand'
)
```

Callable that returns a list of PIL frames for a given file path.

```python
nemo_automodel.components.datasets.multimodal.video.FrameSampler.__call__(
    file_name
)
```

```python
class nemo_automodel.components.datasets.multimodal.transforms.ImageTransform(
    max_image_size,
    min_image_size,
    image_stride,
    max_pixels = 14 * 14 * 9 * 1024,
    image_mean = (0.5, 0.5, 0.5),
    image_std = (0.5, 0.5, 0.5)
)
```

Full BAGEL image transform: resize + to\_tensor + normalize.

Used for both ViT input (stride=14) and VAE input (stride=16, via
separate instances). `stride` is exposed as an attribute so the
dataset can compute patch counts without knowing the transform class.

**`normalize_transform`**

---

**`resize_transform`**

---

**`to_tensor_transform`** `= tv_transforms.ToTensor()`

---

```python
nemo_automodel.components.datasets.multimodal.transforms.ImageTransform.__call__(
    img,
    img_num = 1
)
```

```python
class nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset()
```

**Bases:** [DistributedIterableDataset](/nemo-automodel/nemo_automodel/components/datasets/multimodal/distributed_iterable#nemo_automodel-components-datasets-multimodal-distributed_iterable-DistributedIterableDataset)

Builder mixin for interleaved image/text/video sequence plans.

Subclasses still provide `__init__` + `__iter__` + `parse_row`
(via `ParquetStandardIterableDataset`); this class only holds
the per-item append helpers used inside `parse_row`.

```python
nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_image(
    data,
    image,
    need_loss,
    need_vae,
    need_vit,
    enable_cfg = True
)
```

```python
nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_text(
    data,
    text,
    need_loss,
    enable_cfg = True
)
```

```python
nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._add_video(
    data,
    frames,
    frame_indexes,
    need_loss,
    need_vae,
    enable_cfg = True
)
```

```python
nemo_automodel.components.datasets.multimodal.interleave.InterleavedBaseIterableDataset._init_data()
```

```python
class nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize(
    max_size: int,
    min_size: int,
    stride: int,
    max_pixels: int,
    interpolation = None,
    antialias = True
)
```

**Bases:** `Module`

Resize so longest/shortest edges stay within bounds and both edges are stride-divisible.

**Parameters:**

**`max_size`** `int`

Maximum size for the longest edge.

---

**`min_size`** `int`

Minimum size for the shortest edge.

---

**`stride`** `int`

Value both edges must be divisible by (ViT patch size).

---

**`max_pixels`** `int`

Maximum total pixels for the full image.

---

**`interpolation`** — default: None

Torchvision interpolation mode (default bicubic).

---

**`antialias`** — default: True

Whether to apply antialiasing.

---

```python
nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize._apply_scale(
    width,
    height,
    scale
)
```

```python
nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize._make_divisible(
    value,
    stride
)
```

```python
nemo_automodel.components.datasets.multimodal.transforms.MaxLongEdgeMinShortEdgeResize.forward(
    img,
    img_num = 1
)
```

```python
class nemo_automodel.components.datasets.multimodal.packing.PackedDataset(
    data_config: 'BagelDatasetConfig',
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, object],
    local_rank: int,
    world_size: int,
    num_workers: int,
    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: object | None = None,
    dataset_info: collections.abc.Mapping[str, object] | None = None,
    global_seed: int | None = None
)
```

**Bases:** `IterableDataset`

Greedy pack of samples drawn from weighted groups into token-budgeted batches.

The dataset reseeds at iterator start so AM sees a deterministic
BAGEL-compatible packed-data stream regardless of earlier RNG consumption
during model construction or checkpoint loading.

**`_drop_counters`** `= {}`

---

**`_resume_buffer`** `= []`

---

**`_resume_sequence_status`** `= self.set_sequence_status()`

---

**`_yielded_batches`** `= 0`

---

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.__iter__()
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._grouped_dataset_state_dicts()
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._load_grouped_dataset_state_dicts(
    states
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._load_rng_state_dict(
    state
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._log_drop(
    reason,
    message,
    args = (),
    every = 100
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._rng_state_dict()
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset._set_resume_point(
    buffer,
    yielded_batches
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.build_datasets(
    datasets_metainfo,
    data_status
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.load_state_dict(
    state_dict
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.pack_sequence(
    sample,
    sequence_status
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.set_epoch(
    seed
)
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.set_sequence_status()
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.state_dict()
```

```python
nemo_automodel.components.datasets.multimodal.packing.PackedDataset.to_tensor(
    sequence_status
)
```

```python
class nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset(
    dataset_name,
    transform,
    tokenizer,
    vit_transform,
    data_dir_list,
    num_used_data,
    parquet_info,
    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)

Base class: iterate per-(file, row\_group) across a list of parquet shards.

Subclasses override `parse_row` to turn one pandas row into the
dict schema consumed by `.packing.PackedDataset`.

**`data_paths`**

---

```python
nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.__iter__()
```

```python
nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.get_data_paths(
    data_dir_list,
    num_used_data,
    parquet_info
)
```

```python
nemo_automodel.components.datasets.multimodal.interleave.ParquetStandardIterableDataset.parse_row(
    row
)
```

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

**`data_paths`**

---

```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.collate_fns.SimpleCustomBatch(
    batch
)
```

Pass-through wrapper around one packed batch from `PackedDataset`.

**`attn_modes`** `= data['attn_modes']`

---

**`batch_data_indexes`** `= data['batch_data_indexes']`

---

**`ce_loss_indexes`** `= data['ce_loss_indexes']`

---

**`ce_loss_weights`** `= data['ce_loss_weights']`

---

**`mse_loss_indexes`** `= data['mse_loss_indexes']`

---

**`nested_attention_masks`** `= data['nested_attention_masks']`

---

**`packed_label_ids`** `= data['packed_label_ids']`

---

**`packed_latent_position_ids`** `= data['packed_latent_position_ids']`

---

**`packed_position_ids`** `= data['packed_position_ids']`

---

**`packed_text_ids`** `= data['packed_text_ids']`

---

**`packed_text_indexes`** `= data['packed_text_indexes']`

---

**`packed_timesteps`** `= data['packed_timesteps']`

---

**`packed_vae_token_indexes`** `= data['packed_vae_token_indexes']`

---

**`packed_vit_position_ids`** `= data['packed_vit_position_ids']`

---

**`packed_vit_token_indexes`** `= data['packed_vit_token_indexes']`

---

**`packed_vit_tokens`** `= data['packed_vit_tokens']`

---

**`padded_images`** `= data['padded_images']`

---

**`patchified_vae_latent_shapes`** `= data['patchified_vae_latent_shapes']`

---

**`sample_lens`** `= data['sample_lens']`

---

**`sequence_length`** `= data['sequence_length']`

---

**`split_lens`** `= data['split_lens']`

---

**`use_flex`** `= 'nested_attention_masks' not in data.keys()`

---

**`vit_token_seqlens`** `= data['vit_token_seqlens']`

---

```python
nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.cuda(
    device
)
```

```python
nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.pin_memory()
```

```python
nemo_automodel.components.datasets.multimodal.collate_fns.SimpleCustomBatch.to_dict()
```

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

Each parquet row is expected to carry:

* `image`: raw image bytes.
* `captions`: JSON-encoded `{key: caption}` dict; one caption is
  sampled uniformly per row.

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

---

```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
class nemo_automodel.components.datasets.multimodal.interleave.UnifiedEditIterableDataset()
```

**Bases:** [InterleavedBaseIterableDataset](#nemo_automodel-components-datasets-multimodal-interleave-InterleavedBaseIterableDataset), [ParquetStandardIterableDataset](#nemo_automodel-components-datasets-multimodal-interleave-ParquetStandardIterableDataset)

Image-editing dataset: `(input, instruction, output)` chains over parquet.

Row schema (upstream BAGEL `seedxedit_multi` + compatibles):
`image_list`: list of raw image bytes (at least 2).
`instruction_list`: list of lists; `instruction_list[i]` is a
set of equivalent phrasings for the edit that turns
`image_list[i]` into `image_list[i+1]`.

```python
nemo_automodel.components.datasets.multimodal.interleave.UnifiedEditIterableDataset.parse_row(
    row
)
```

```python
nemo_automodel.components.datasets.multimodal.utils.add_special_tokens(
    tokenizer
)
```

Add BAGEL's four special tokens to a Qwen2 tokenizer and return the ids.

**Returns:**

tuple `(tokenizer, new_token_ids, num_new_tokens)` where

```python
nemo_automodel.components.datasets.multimodal.collate_fns.bagel_packed_collate_fn(
    batch
)
```

Canonical name in AM's collate-fn registry.

```python
nemo_automodel.components.datasets.multimodal.collate_fns.collate_wrapper()
```

Return the BAGEL-style identity collate (wraps a single packed dict).

```python
nemo_automodel.components.datasets.multimodal.utils.get_flattened_position_ids_extrapolate(
    img_h,
    img_w,
    patch_size,
    max_num_patches_per_side
)
```

Return flattened 2D patch position IDs by direct grid extrapolation.

```python
nemo_automodel.components.datasets.multimodal.utils.get_flattened_position_ids_interpolate(
    img_h,
    img_w,
    patch_size,
    max_num_patches_per_side
)
```

Return flattened 2D patch position IDs by interpolating to the max grid.

```python
nemo_automodel.components.datasets.multimodal.parquet_utils.get_parquet_data_paths(
    data_dir_list,
    num_sampled_data_paths,
    rank = 0,
    world_size = 1
)
```

Return a flat list of parquet file paths sharded across ranks.

Directories are split across ranks via
`chunk_size = ceil(num_dirs / world_size)`. Each rank lists its local
directories, repeats the file list to reach `num_sampled_data_paths` per
directory, then all-gathers across ranks so every rank ends up with the
same combined list.

```python
nemo_automodel.components.datasets.multimodal.parquet_utils.init_arrow_pf_fs(
    parquet_file_path
)
```

Return a pyarrow filesystem for `parquet_file_path`.

`pyarrow` is imported lazily because not every AM install carries it,
and the import is only needed when an actually-parquet-backed dataset
(T2I, UnifiedEdit) is constructed.

```python
nemo_automodel.components.datasets.multimodal.utils.len2weight(
    x,
    loss_reduction = 'square'
)
```

Convert a sequence length into BAGEL's per-sample loss weight.

```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 `PackedDataset` instance. The caller is
still responsible for invoking `set_epoch` immediately before
DataLoader iteration (and for the pre-iter RNG reseed — see
`PackedDataset.set_epoch` docstring).

```python
nemo_automodel.components.datasets.multimodal.utils.patchify(
    image,
    patch_size
)
```

Patchify a CxHxW image tensor into (H/p \* W/p, p*p*C).

```python
nemo_automodel.components.datasets.multimodal.utils.pil_img2rgb(
    image
)
```

Convert a PIL image to RGB, compositing transparent pixels on white.

```python
nemo_automodel.components.datasets.multimodal.utils.prepare_attention_mask_per_sample(
    split_lens,
    attn_modes,
    device = 'cpu'
)
```

Build a per-sample additive float mask honoring causal / full / noise splits.

```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.packing.DataConfig = BagelDatasetConfig
```