nemo_automodel.components.datasets.lazy_mapped_dataset

Module Contents

Classes

Name	Description
`LazyMappedDataset`	Dataset wrapper that applies a transform function on-the-fly instead of
`LazyMappedDatasetConfig`	Construction-time configuration for :class:`LazyMappedDataset`.

Data

logger

API

class nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset(
    dataset,
    map_fn,
    cache_size = 10000
)

Bases: Dataset

Dataset wrapper that applies a transform function on-the-fly instead of preprocessing the whole dataset upfront with .map(fn).

Parameters:

dataset

Any object that supports __len__ and __getitem__ (e.g. a Hugging Face datasets.Dataset).

map_fn

A callable that accepts a single example and returns the transformed example.

cache_size

Defaults to 10000

Number of processed items to cache. Defaults to the 10k dataset samples. Set to 0 to disable caching or None to cache all.

Returns:

A map-style dataset that applies map_fn lazily on each item access.

cache_info

Any | None

Return LRU cache statistics, or None if caching is disabled.

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset.__getitem__(
    idx: int
) -> typing.Any

Returns the transformed item at the given index

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset.__getstate__() -> dict

Returns pickable state by dropping the unpicklable _get_item function

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset.__len__() -> int

Returns the number of items in the dataset

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset.__repr__() -> str

returns a string representation of the dataset

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset.__setstate__(
    state: dict
) -> None

Restores state and rebuild _get_item after unpickling

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset._build_get_item() -> None

Build the internal item accessor, with or without LRU caching

class nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDatasetConfig(
    cache_size: int | None = 10000
)

Dataclass

Construction-time configuration for :class:LazyMappedDataset.

dataset and map_fn are runtime objects rather than serializable config, so they are explicit :meth:build arguments.

cache_size

int | None = 10000

Number of processed items to cache (0 disables caching, None caches all).

nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDatasetConfig.build(
    dataset: torch.utils.data.Dataset,
    map_fn: collections.abc.Callable[[object], object]
) -> 'LazyMappedDataset'

Build a lazy mapping wrapper around a runtime dataset and transform.

Parameters:

dataset

Dataset

Map-style dataset whose samples are transformed lazily.

map_fn

Callable[[object], object]

Callable that converts one source sample into one output sample.

Returns: 'LazyMappedDataset'

Lazy mapped dataset with the configured cache size.

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