dockyard_rl.data.dataloader

Module Contents

Classes

MultipleDataloaderWrapper

Wrapper for multiple dataloaders.

API

class dockyard_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/configs/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

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.