nemo_rl.data.dataloader#
Module Contents#
Classes#
Wrapper for multiple dataloaders. |
API#
- 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.