nemo_automodel.components.datasets.llm.mock_iterable_dataset

Module Contents

Classes

Name	Description
`MockIterableDataset`	Mock dataset that generates synthetic data for benchmarking.
`MockIterableDatasetConfig`	Construction-time configuration for :class:`MockIterableDataset`.

API

class nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDataset(
    vocab_size: int = 1024,
    seq_len: int = 1024,
    num_samples: int = 1000000,
    batch_size: int = 1
)

Bases: IterableDataset

Mock dataset that generates synthetic data for benchmarking.

This dataset generates random tokens similar to the benchmarking script, creating input_ids, labels, and position_ids for each sample.

nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDataset.__iter__()

Generate synthetic batches.

nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDataset.__len__()

Return the number of samples.

class nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDatasetConfig(
    vocab_size: int = 1024,
    seq_len: int = 1024,
    num_samples: int = 1000000,
    batch_size: int = 1
)

Dataclass

Construction-time configuration for :class:MockIterableDataset.

batch_size

int = 1

Batch size to yield (1 for unbatched samples).

num_samples

int = 1000000

Total number of samples to generate (1M for an infinite-like dataset).

seq_len

int = 1024

Sequence length for each sample.

vocab_size

int = 1024

Size of the vocabulary for generating random tokens.

nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDatasetConfig.build() -> 'MockIterableDataset'

Build a :class:MockIterableDataset from this :class:MockIterableDatasetConfig.