`nemo_rl.data.dataloader`#

Module Contents#

Classes#

`CyclingDataLoader`	Repeat a stateful dataloader until its consumer stops.
`MultipleDataloaderWrapper`	Wrapper for multiple dataloaders.

API#

class nemo_rl.data.dataloader.CyclingDataLoader( dataloader: torchdata.stateful_dataloader.StatefulDataLoader, )#

Repeat a stateful dataloader until its consumer stops.

Initialization

__iter__() → collections.abc.Iterator[nemo_rl.distributed.batched_data_dict.BatchedDataDict]#

state_dict() → dict[str, Any]#

class nemo_rl.data.dataloader.MultipleDataloaderWrapper( expected_num_prompts: int, data_config: dict, dataloaders: dict[str, torchdata.stateful_dataloader.StatefulDataLoader], )#

Wrapper for multiple dataloaders.

This wrapper is used to sample data from multiple dataloaders using a custom dataloader function.

When a single dataloader is exhausted, the data iterator must be reset in the custom dataloader function (as demonstrated in examples/custom_dataloader/custom_dataloader.py). This design ensures that the MultipleDataloaderWrapper operates as an infinite iterator, where next() will not raise StopIteration and len() is not supported.

Initialization

_load_custom_dataloader_func()#

__iter__()#

__next__()#

set_records(records: dict)#

Set the records for the custom dataloader.

Records are used to pass additional information to the custom dataloader function to decide how to sample the data from the dataloaders.

nemo_rl.data.dataloader#

Module Contents#

Classes#

API#

`nemo_rl.data.dataloader`#