nemo_automodel.components.datasets.multimodal.loader

Typed construction for the packed BAGEL multimodal dataloader.

Module Contents

Classes

Name	Description
`BagelDataloaderBuild`	Materialized BAGEL dataloader and its packed dataset.
`BagelDataloaderConfig`	Construction-time configuration for the complete BAGEL input pipeline.

Data

__all__

API

class nemo_automodel.components.datasets.multimodal.loader.BagelDataloaderBuild(
    dataloader: torchdata.stateful_dataloader.StatefulDataLoader,
    dataset: nemo_automodel.components.datasets.multimodal.packing.PackedDataset
)

Dataclass

Materialized BAGEL dataloader and its packed dataset.

dataloader

StatefulDataLoader

dataset

PackedDataset

class nemo_automodel.components.datasets.multimodal.loader.BagelDataloaderConfig(
    dataset_config: nemo_automodel.components.datasets.multimodal.datasets.BagelDatasetConfig,
    num_workers: int = 1,
    pin_memory: bool = True,
    prefetch_factor: int = 2
)

Dataclass

Construction-time configuration for the complete BAGEL input pipeline.

dataset_config

BagelDatasetConfig

num_workers

int = 1

pin_memory

bool = True

prefetch_factor

int = 2

nemo_automodel.components.datasets.multimodal.loader.BagelDataloaderConfig.build(
    tokenizer: 'PreTrainedTokenizerBase',
    special_tokens: collections.abc.Mapping[str, object],
    rank: int,
    world_size: int,
    batch_size: int,
    global_seed: int
) -> nemo_automodel.components.datasets.multimodal.loader.BagelDataloaderBuild

Build the packed BAGEL dataset and its stateful dataloader.

Parameters:

tokenizer

'PreTrainedTokenizerBase'

Runtime tokenizer used by the packed dataset.

special_tokens

Mapping[str, object]

Runtime BAGEL special-token mapping.

rank

int

Global distributed rank used for source-data sharding and worker seeding.

world_size

int

Global distributed world size.

batch_size

int

Runtime local batch size; BAGEL packed samples require one row per step.

global_seed

int

Global training seed used to initialize dataloader worker RNG state.

Returns: BagelDataloaderBuild

Named result containing the dataloader and its packed dataset.

nemo_automodel.components.datasets.multimodal.loader.__all__ = ['BagelDataloaderBuild', 'BagelDataloaderConfig']