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,
    exclude_token_ids: tuple[int, ...] | list[int] | None = None
)

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.

_allowed_ids

Tensor | None = None

exclude_token_ids

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.

nemo_automodel.components.datasets.llm.mock_iterable_dataset.MockIterableDataset._sample_tokens() -> torch.Tensor

Uniform random [batch_size, seq_len] ids over the vocabulary minus exclude_token_ids.

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,
    exclude_token_ids: tuple[int, ...] | list[int] | None = None
)

Dataclass

Construction-time configuration for :class:MockIterableDataset.

batch_size

int = 1

Batch size to yield (1 for unbatched samples).

exclude_token_ids

tuple[int, ...] | list[int] | None = None

Token ids never emitted in input_ids (e.g. the model’s pad_token_id).

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.