ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDatasetsVlm

nemo_automodel.components.datasets.vlm.datasets

View as Markdown

Module Contents

Classes

NameDescription
PreTokenizedDatasetWrapperDataset wrapper that tokenizes samples in __getitem__.
RobustDatasetWrapperWrapper that catches __getitem__ and collate errors, substituting random replacement samples.
_ExamplesWithStatslist subclass that carries pre-computed dataset statistics.

Functions

NameDescription
_collect_sample_statsCount images, videos, text-only samples and estimate token counts.
_convert_sharegpt_to_conversationConvert a single sharegpt-format example to Automodel conversation format.
_grid_media_token_count-
_load_json_or_jsonlLoad data from a JSON or JSONL file.
_load_jsonl_for_rankLoad only the JSONL lines needed for this rank, avoiding full json.loads on skipped lines.
_log_dataset_loading_summaryPrint a visual summary of per-dataset loading times and data statistics.
_media_token_mismatchReturn a mismatch description if media grids survived without tokens.
_resolve_processor_token_idResolve a model-specific media token id from processor/config/tokenizer.
make_cord_v2_datasetLoad and preprocess the CORD-V2 dataset for image-to-text fine-tuning.
make_llava_onevision_datasetLoad and preprocess the LLaVA-Instruct-150K dataset for LLaVA-OneVision-1.5.
make_medpix_datasetLoad and preprocess the MedPix dataset for image-to-text fine-tuning.
make_meta_datasetLoad datasets defined in a meta JSON file and convert to Automodel conversation format.
make_rdr_datasetLoad and preprocess the RDR dataset for image-to-text fine-tuning.
make_tulu3_datasetLoad allenai/tulu-3-sft-mixture directly from the HF Hub as text-only conversations.
make_tulu3_magicoder_text_mix_datasetBuild a text-only 80/20 mix of Tulu-3-SFT-mixture and Magicoder-OSS-Instruct-75K.
make_unimm_chat_datasetLoad and preprocess the UniMM-Chat dataset for image-to-text fine-tuning.

Data

logger

API

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.

class nemo_automodel.components.datasets.vlm.datasets._ExamplesWithStats()

Bases: list

list subclass that carries pre-computed dataset statistics.

Attached by :func:make_meta_dataset so downstream code (e.g. _log_global_dataset_stats) can read aggregated stats without re-scanning all examples.

__slots__
= ('stats',)
nemo_automodel.components.datasets.vlm.datasets._collect_sample_stats(
    examples
)

Count images, videos, text-only samples and estimate token counts.

Token estimation mirrors the logic in LengthGroupedSampler._estimate_tokens:

  • Text tokens: uses pre-computed _text_tokens when present (written by scripts/precompute_tokens.py), otherwise falls back to chars // 3.
  • Media tokens: uses mm_inputs_meta image/video dimensions when present (populated by the precompute script), otherwise 500 per media item.

Returns:

dict with keys n_images, n_videos, n_text_only, n_text_tokens,

nemo_automodel.components.datasets.vlm.datasets._convert_sharegpt_to_conversation(
    example,
    columns = None,
    tags = None,
    media_dir = None
)

Convert a single sharegpt-format example to Automodel conversation format.

Parameters:

example
dict

A single data example in sharegpt format.

columns
dictDefaults to None

Column name mapping with keys ‘messages’, ‘images’, ‘videos’.

tags
dictDefaults to None

Tag mapping with keys ‘role_tag’, ‘content_tag’, ‘user_tag’, ‘assistant_tag’.

media_dir
str | NoneDefaults to None

Directory prefix for resolving relative media paths.

Returns:

Example in Automodel conversation format.

nemo_automodel.components.datasets.vlm.datasets._grid_media_token_count(
    grid,
    merge_size: int
) -> int
 
nemo_automodel.components.datasets.vlm.datasets._load_json_or_jsonl(
    file_path
)

Load data from a JSON or JSONL file.

Parameters:

file_path
str

Path to the JSON or JSONL file.

Returns:

list[dict]: List of data examples.

nemo_automodel.components.datasets.vlm.datasets._load_jsonl_for_rank(
    file_path,
    sample_ratio,
    rank,
    world_size
)

Load only the JSONL lines needed for this rank, avoiding full json.loads on skipped lines.

Handles sample_ratio and sharding so that each rank only parses and stores its own subset. The semantics match the original load-all-then-slice approach:

  1. Apply sample_ratio (deterministic Random(42).sample) on the full index range.
  2. Shard the resulting list with [rank::world_size].

Returns:

tuple[list[dict], int]: (parsed examples for this rank, total line count).

nemo_automodel.components.datasets.vlm.datasets._log_dataset_loading_summary(
    timings,
    wall_time,
    total_samples,
    rank = None
)

Print a visual summary of per-dataset loading times and data statistics.

nemo_automodel.components.datasets.vlm.datasets._media_token_mismatch(
    input_ids,
    result,
    processor
) -> str | None

Return a mismatch description if media grids survived without tokens.

nemo_automodel.components.datasets.vlm.datasets._resolve_processor_token_id(
    processor,
    attr_names,
    token_names
)

Resolve a model-specific media token id from processor/config/tokenizer.

nemo_automodel.components.datasets.vlm.datasets.make_cord_v2_dataset(
    path_or_dataset = 'naver-clova-ix/cord-v2',
    split = 'train',
    kwargs = {}
)

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

nemo_automodel.components.datasets.vlm.datasets.make_llava_onevision_dataset(
    path_or_dataset = 'liuhaotian/LLaVA-Instruct-...,
    split = 'train',
    kwargs = {}
)

Load and preprocess the LLaVA-Instruct-150K dataset for LLaVA-OneVision-1.5.

This function loads conversation-format data with images and returns it in the standard NeMo VLM format expected by the collate function.

Parameters:

path_or_dataset
Defaults to 'liuhaotian/LLaVA-Instruct-150K'

Path to the dataset on HuggingFace Hub or local path.

split
Defaults to 'train'

Dataset split to load (e.g., “train”, “train[:1000]”).

**kwargs
Defaults to {}

Additional arguments passed to load_dataset.

Returns:

List of dicts with “conversation” and “image” keys.

nemo_automodel.components.datasets.vlm.datasets.make_medpix_dataset(
    path_or_dataset = 'medpix-dataset/medpix-data...,
    split = 'train',
    kwargs = {}
)

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

Formatting is deferred to __getitem__ via with_transform so the dataset stays Arrow-backed (no Python-side copy of every image). Images are loaded undecoded (Image(decode=False)) and wrapped as lazy PIL handles (Image.open reads only the header); the actual pixel decode then happens on demand in the DataLoader workers, rather than eagerly decoding the whole split up front.

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).
  • 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_tulu3_dataset(
    path_or_dataset: str = 'allenai/tulu-3-sft-mixture',
    split: str = 'train',
    kwargs = {}
)

Load allenai/tulu-3-sft-mixture directly from the HF Hub as text-only conversations.

This avoids the meta-JSON + JSONL data-prep step required by :func:make_meta_dataset: point the recipe’s dataset._target_ at this function and the Tulu-3 split is pulled straight from the Hub. Each row’s messages field is converted with the same helper the meta-JSON path uses (:func:_convert_sharegpt_to_conversation), so the resulting data composition is identical to dumping the split to JSONL and loading it via :func:make_meta_dataset: no turn cap, system turns dropped, every row kept in the original split order. Conversations are text-only (no image entries), so batches carry no pixel_values / vision tensors.

The returned dataset stays Arrow-backed (map) so the full ~939k-row split is not copied into a Python list.

Parameters:

path_or_dataset
strDefaults to 'allenai/tulu-3-sft-mixture'

HF Hub id (or local path) of the Tulu-3 SFT mixture.

split
strDefaults to 'train'

HF split expression (e.g. "train" or "train[:50000]").

**kwargs
Defaults to {}

Ignored. Accepted so recipe-level dataset keys (e.g. truncate) that are forwarded to the dataset target do not raise.

Returns:

datasets.Dataset: Rows with a single conversation column, each a list of

nemo_automodel.components.datasets.vlm.datasets.make_tulu3_magicoder_text_mix_dataset(
    tulu_split: str = 'train',
    magicoder_split: str = 'train',
    seed: int = 42,
    max_turns: int = 16,
    limit_total: int | None = None,
    kwargs = {}
) -> list

Build a text-only 80/20 mix of Tulu-3-SFT-mixture and Magicoder-OSS-Instruct-75K.

Both datasets are converted into the NeMo VLM {"conversation": [...]} shape consumed by :func:nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn. Because default_collate_fn is image-aware only when a conversation turn contains an {"type": "image", ...} entry, returning text-only conversations here yields batches with no pixel_values / vision tensors — which is what the Gemma 4 base+drafter composite expects for text-only training.

Mixing uses datasets.interleave_datasets with probabilities [0.8, 0.2] and stopping_strategy="all_exhausted" so both datasets are sampled until every example has been drawn at least once.

Parameters:

tulu_split
strDefaults to 'train'

HF split expression for the Tulu-3 source (e.g. "train" or "train[:50000]").

magicoder_split
strDefaults to 'train'

HF split expression for the Magicoder source.

seed
intDefaults to 42

Seed forwarded to interleave_datasets for reproducibility.

max_turns
intDefaults to 16

Drop Tulu-3 conversations with more than this many turns to keep memory bounded. Magicoder samples are always 2 turns.

limit_total
int | NoneDefaults to None

If set, cap the merged dataset to this many rows.

**kwargs
Defaults to {}

Additional arguments forwarded to load_dataset for both sources.

Returns: list

List of {"conversation": [...]} dicts. Each conversation is a list of

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.

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