nemo_automodel.components.datasets.vlm

Submodules

Package Contents

Classes

Name	Description
`LengthGroupedSampler`	Sampler that groups samples by total token count for balanced
`PreTokenizedDatasetWrapper`	Dataset wrapper that tokenizes samples in `__getitem__`.
`RobustDatasetWrapper`	Wrapper that catches `__getitem__` and collate errors, substituting random replacement samples.

Functions

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

API

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

Defaults to 42

Base random seed (same value on every rank).

processor

Defaults to 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

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

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

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

Compute token lengths with direct list access for speed.

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

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

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.

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.

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

staticmethod

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

staticmethod

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

staticmethod

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

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

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

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

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

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

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

Wrap collate_fn so that on failure the entire batch is re-sampled.

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.

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

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

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.

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

intDefaults to 10

Number of conversation examples to generate.

num_images_per_sample

intDefaults to 1

Number of random images per user turn.

image_size

Tuple[int, int]Defaults to (256, 256)

(width, height) of each generated image.

prompt

strDefaults to 'Describe this image.'

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

responses

List[str] | NoneDefaults to None

Optional list of assistant responses. Cycled over samples.

max_length

int | NoneDefaults to 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

intDefaults to 0

Random seed for reproducibility.

Returns: list

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

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.

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] | NoneDefaults to None

Which datasets to load. None means all.

split

strDefaults to 'train'

Unused, kept for API consistency.

shard_data

boolDefaults to False

If True, each rank loads only its 1/world_size slice.

rank

int | NoneDefaults to None

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

world_size

int | NoneDefaults to None

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

**kwargs

Defaults to {}

Additional arguments (unused).

Returns:

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

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

strDefaults to 'quintend/rdr-items'

Path or identifier for the RDR dataset.

split

strDefaults to 'train'

Dataset split to load.

**kwargs

Defaults to {}

Additional arguments.

Returns:

The processed dataset.

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.