nemo_rl.data.utils#

Module Contents#

Functions#

setup_data_with_envs

Setup data with environments.

API#

nemo_rl.data.utils.setup_data_with_envs(
tokenizer: transformers.AutoProcessor | transformers.AutoTokenizer,
data_config: nemo_rl.data.DataConfig,
env_configs: dict[str, Any],
is_vlm: bool = False,
) → tuple[nemo_rl.data.datasets.AllTaskProcessedDataset, Optional[nemo_rl.data.datasets.AllTaskProcessedDataset], dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]]#

Setup data with environments.

This function is used to setup the data and environments for the training and validation datasets.

Parameters:

  • tokenizer – Tokenizer or processor.

  • data_config – Data config.

  • env_configs – Environment configs.

  • is_vlm – Whether to use VLM training or not.

Returns:

A tuple of (train dataset, validation dataset, task to environment, task to validation environment).