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

## Submodules

* **[`nemo_automodel.components.datasets.vlm.collate_fns`](/nemo-automodel/nemo_automodel/components/datasets/vlm/collate_fns)**
* **[`nemo_automodel.components.datasets.vlm.datasets`](/nemo-automodel/nemo_automodel/components/datasets/vlm/datasets)**
* **[`nemo_automodel.components.datasets.vlm.dspark_collate`](/nemo-automodel/nemo_automodel/components/datasets/vlm/dspark_collate)**
* **[`nemo_automodel.components.datasets.vlm.fake_image`](/nemo-automodel/nemo_automodel/components/datasets/vlm/fake_image)**
* **[`nemo_automodel.components.datasets.vlm.loader`](/nemo-automodel/nemo_automodel/components/datasets/vlm/loader)**
* **[`nemo_automodel.components.datasets.vlm.mock`](/nemo-automodel/nemo_automodel/components/datasets/vlm/mock)**
* **[`nemo_automodel.components.datasets.vlm.neat_packing_vlm`](/nemo-automodel/nemo_automodel/components/datasets/vlm/neat_packing_vlm)**
* **[`nemo_automodel.components.datasets.vlm.pp_media`](/nemo-automodel/nemo_automodel/components/datasets/vlm/pp_media)**
* **[`nemo_automodel.components.datasets.vlm.samplers`](/nemo-automodel/nemo_automodel/components/datasets/vlm/samplers)**
* **[`nemo_automodel.components.datasets.vlm.utils`](/nemo-automodel/nemo_automodel/components/datasets/vlm/utils)**

## Package Contents

### Classes

| Name                                                                                                        | Description                                                                                     |
| ----------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- |
| [`LengthGroupedSampler`](#nemo_automodel-components-datasets-vlm-samplers-LengthGroupedSampler)             | Sampler that groups samples by total token count for balanced                                   |
| [`PreTokenizedDatasetWrapper`](#nemo_automodel-components-datasets-vlm-datasets-PreTokenizedDatasetWrapper) | Dataset wrapper that tokenizes samples in `__getitem__`.                                        |
| [`RobustDatasetWrapper`](#nemo_automodel-components-datasets-vlm-datasets-RobustDatasetWrapper)             | Wrapper that catches `__getitem__` and collate errors, substituting random replacement samples. |

### Functions

| Name                                                                                                  | Description                                                                             |
| ----------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------- |
| [`build_mock_vlm_dataset`](#nemo_automodel-components-datasets-vlm-mock-build_mock_vlm_dataset)       | Build a mock VLM dataset in Automodel conversation format.                              |
| [`make_cord_v2_dataset`](#nemo_automodel-components-datasets-vlm-datasets-make_cord_v2_dataset)       | Load and preprocess the CORD-V2 dataset for image-to-text fine-tuning.                  |
| [`make_meta_dataset`](#nemo_automodel-components-datasets-vlm-datasets-make_meta_dataset)             | Load datasets defined in a meta JSON file and convert to Automodel conversation format. |
| [`make_rdr_dataset`](#nemo_automodel-components-datasets-vlm-datasets-make_rdr_dataset)               | Load and preprocess the RDR dataset for image-to-text fine-tuning.                      |
| [`make_unimm_chat_dataset`](#nemo_automodel-components-datasets-vlm-datasets-make_unimm_chat_dataset) | Load and preprocess the UniMM-Chat dataset for image-to-text fine-tuning.               |

### API

```python
class nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler(
    dataset,
    seed = 42,
    processor = None,
    max_length = None,
    batch_size = 1
)
```

**Bases:** `Sampler`

Sampler that groups samples by total token count for balanced
distributed training.

With `shard_data=True` each rank owns a different subset of data.
This sampler sorts every rank's indices by **total tokens**
(`text_tokens + media_tokens`, descending).  All ranks share the
same `seed + epoch` so position *N* on every rank corresponds to a
sample of similar length, keeping cross-rank padding minimal.

Per-epoch randomness is achieved by rotating the sorted order by a
deterministic random offset (same on every rank).

**Parameters:**

**`dataset`**

The dataset to sample from.

---

**`seed`** — default: 42

Base random seed (same value on every rank).

---

**`processor`** — default: None

Optional HuggingFace processor (e.g. `Qwen2VLProcessor`).
Used to read `image_processor` / `video_processor` attributes
for accurate media token estimation via `smart_resize`.

---

**`_image_cfg`**

---

**`_video_cfg`**

---

**`batch_size`** `= max(1, batch_size)`

---

**`epoch`** `= 0`

---

**`lengths`**

---

**`sorted_indices`**

---

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler.__iter__()
```

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler.__len__()
```

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._compute_or_load_lengths(
    dataset
)
```

Compute token lengths with direct list access for speed.

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._estimate_image_tokens(
    img_meta
)
```

Estimate token count for one image from its `[height, width]` metadata.

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._estimate_tokens(
    example
)
```

Return `(text_tokens, media_tokens)` for one example.

Uses pre-computed `_text_tokens` / `_media_tokens` when available
(written by `scripts/precompute_tokens.py`).  Otherwise falls back
to heuristic estimation.

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._estimate_video_tokens(
    vid_meta
)
```

Estimate token count for one video from its
`[total_frames, height, width, fps, duration]` metadata.

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._extract_image_config(
    processor
)
```

staticmethod

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._extract_video_config(
    processor
)
```

staticmethod

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler._get_raw_samples(
    dataset
)
```

staticmethod

Unwrap dataset wrappers to get the underlying list for direct access.

```python
nemo_automodel.components.datasets.vlm.samplers.LengthGroupedSampler.set_epoch(
    epoch
)
```

Set the epoch for deterministic shuffling (standard PyTorch pattern).

```python
class nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapper(
    dataset,
    processor,
    max_length = None,
    max_retries = 10,
    truncate = False,
    post_tokenize_hook = None,
    inject_fake_images = True
)
```

**Bases:** `Dataset`

Dataset wrapper that tokenizes samples in `__getitem__`.

Instead of deferring `apply_chat_template` to the collate function, this
wrapper performs tokenization per-sample so that:

* The collate function only needs to pad and stack.
* Overlong samples are detected **after** precise tokenization (including
  media-token expansion) and replaced with a different random sample.
* Tokenization work is distributed across DataLoader workers.

Each `__getitem__` call returns a dict with at least::

\{
"input\_ids":      (seq\_len,),
"attention\_mask": (seq\_len,),
"labels":         (seq\_len,),
}

Plus optional media tensors (`pixel_values`, `image_grid_thw`,
`pixel_values_videos`, `video_grid_thw`).

```python
nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapper.__getitem__(
    idx
)
```

```python
nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapper.__len__()
```

```python
nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapper.robust_collate(
    collate_fn
)
```

Wrap *collate\_fn* so that on failure the entire batch is re-sampled.

```python
class nemo_automodel.components.datasets.vlm.datasets.RobustDatasetWrapper(
    dataset,
    max_retries: int = 10
)
```

**Bases:** `Dataset`

Wrapper that catches `__getitem__` and collate errors, substituting random replacement samples.

This handles failures such as corrupted files, missing media, bad data,
or processor errors (e.g. multimodal token mismatch from truncation)
without crashing the entire training run.

```python
nemo_automodel.components.datasets.vlm.datasets.RobustDatasetWrapper.__getitem__(
    idx
)
```

```python
nemo_automodel.components.datasets.vlm.datasets.RobustDatasetWrapper.__len__()
```

```python
nemo_automodel.components.datasets.vlm.datasets.RobustDatasetWrapper.robust_collate(
    collate_fn
)
```

Wrap a collate\_fn so that on failure the entire batch is re-sampled and retried.

```python
nemo_automodel.components.datasets.vlm.mock.build_mock_vlm_dataset(
    num_samples: int = 10,
    num_images_per_sample: int = 1,
    image_size: typing.Tuple[int, int] = (256, 256),
    prompt: str = 'Describe this image.',
    responses: typing.List[str] | None = None,
    max_length: int | None = None,
    seed: int = 0,
    kwargs = {}
) -> list
```

Build a mock VLM dataset in Automodel conversation format.

Each sample is a dict with a `"conversation"` key whose value is a list
of user/assistant message dicts.  User messages contain one or more
`{"type": "image", "image": <PIL.Image>}` items followed by a text prompt.
Assistant messages contain a single text response.

This is the same format produced by `make_rdr_dataset`,
`make_unimm_chat_dataset`, and `make_meta_dataset`, so the returned
list can be fed directly to `PreTokenizedDatasetWrapper`.

When `max_length` is set and `responses` is `None`, each sample's
assistant response is generated with `max_length` words — guaranteed
to exceed `max_length` tokens so that `PreTokenizedDatasetWrapper`
with `truncate=True` produces exactly `max_length` tokens per sample.

**Parameters:**

**`num_samples`** `int` — default: 10

Number of conversation examples to generate.

---

**`num_images_per_sample`** `int` — default: 1

Number of random images per user turn.

---

**`image_size`** `Tuple[int, int]` — default: (256, 256)

`(width, height)` of each generated image.

---

**`prompt`** `str` — default: 'Describe this image.'

Text prompt appended after the image(s) in the user turn.

---

**`responses`** `List[str] | None` — default: None

Optional list of assistant responses.  Cycled over samples.

---

**`max_length`** `int | None` — default: None

Target sequence length.  When set (and `responses` is
`None`), generates a response of `max_length` words per sample
so the tokenized sequence always exceeds `max_length` tokens.

---

**`seed`** `int` — default: 0

Random seed for reproducibility.

---

**Returns:** `list`

A list of dicts, each with a single `"conversation"` key.

```python
nemo_automodel.components.datasets.vlm.datasets.make_cord_v2_dataset(
    path_or_dataset = 'naver-clova-ix/cord-v2',
    split = 'train',
    limit_dataset_samples: int | None = None,
    kwargs = {}
)
```

Load and preprocess the CORD-V2 dataset for image-to-text fine-tuning.

```python
nemo_automodel.components.datasets.vlm.datasets.make_meta_dataset(
    path_or_dataset,
    dataset_names = None,
    split = 'train',
    shard_data = False,
    rank = None,
    world_size = None,
    kwargs = {}
)
```

Load datasets defined in a meta JSON file and convert to Automodel conversation format.

The meta JSON file maps dataset names to their configurations. Each configuration can have:

* file\_name (str): Path to the data file (JSON/JSONL). Relative paths are resolved
  against the meta file's directory.
* columns (dict): Column name mapping (messages, images, videos).
* tags (dict): Tag mapping (role\_tag, content\_tag, user\_tag, assistant\_tag, system\_tag).
* media\_dir (str): Directory prefix for media files.
* sample\_ratio (float): Sampling ratio (0.0 to 1.0, default 1.0).

When `shard_data=True`, each rank loads only its `1/world_size` slice of
every dataset file (interleaved: `raw_data[rank::world_size]`).  This
reduces per-rank memory and I/O.  The caller should use a local sampler
(e.g. `RandomSampler`) instead of `DistributedSampler` since data is
already partitioned.

Video frame sampling (fps, min\_frames, max\_frames) should be configured on
the **processor** rather than here.  For example in YAML::

processor:
*target*: transformers.AutoProcessor.from\_pretrained
pretrained\_model\_name\_or\_path: ...
fps: 1
min\_frames: 4
max\_frames: 128

Example meta JSON::

\{
"my\_dataset": \{
"file\_name": "data/train.jsonl",
"columns": \{"messages": "conversations"},
"media\_dir": "/data/media"
}
}

**Parameters:**

**`path_or_dataset`** `str`

Path to the meta JSON file.

---

**`dataset_names`** `list[str] | None` — default: None

Which datasets to load. None means all.

---

**`split`** `str` — default: 'train'

Unused, kept for API consistency.

---

**`shard_data`** `bool` — default: False

If True, each rank loads only its 1/world\_size slice.

---

**`rank`** `int | None` — default: None

Data-parallel rank. Inferred from torch.distributed if None.

---

**`world_size`** `int | None` — default: None

Data-parallel world size. Inferred from torch.distributed if None.

---

**`**kwargs`** — default: \{}

Additional arguments (unused).

---

**Returns:**

list\[dict]: Combined list of examples in Automodel conversation format.

```python
nemo_automodel.components.datasets.vlm.datasets.make_rdr_dataset(
    path_or_dataset = 'quintend/rdr-items',
    split = 'train',
    kwargs = {}
)
```

Load and preprocess the RDR dataset for image-to-text fine-tuning.

**Parameters:**

**`path_or_dataset`** `str` — default: 'quintend/rdr-items'

Path or identifier for the RDR dataset.

---

**`split`** `str` — default: 'train'

Dataset split to load.

---

**`**kwargs`** — default: \{}

Additional arguments.

---

**Returns:**

The processed dataset.

```python
nemo_automodel.components.datasets.vlm.datasets.make_unimm_chat_dataset(
    path_or_dataset = 'Yirany/UniMM-Chat',
    split = 'train',
    kwargs = {}
)
```

Load and preprocess the UniMM-Chat dataset for image-to-text fine-tuning.