nemo_rl.data.datasets.utils#

Module Contents#

Functions#

load_audio_from_file

Decode an audio file (or the audio track of a video) as a 1-D float32 array.

assert_no_double_bos

Assert that there are no double starting BOS tokens in the message.

pil_to_base64

Converts a PIL Image object to a base64 encoded string.

load_dataset_from_path

Load a dataset from a local file, huggingface dataset, or Arrow dataset (saved with save_to_disk).

resolve_external_dataset_class

Resolve a fully-qualified dotted dataset path to a class.

update_single_dataset_config

Fill the single dataset config with default dataset config.

merge_datasets

Merge map-style datasets while supporting non-HF wrappers.

extract_necessary_env_names

Extract the necessary environment names from the data config.

get_huggingface_cache_path

Data#

TokenizerType

API#

nemo_rl.data.datasets.utils.TokenizerType#

None

nemo_rl.data.datasets.utils.load_audio_from_file(
path: str,
sampling_rate: int = 16000,
) → numpy.ndarray#

Decode an audio file (or the audio track of a video) as a 1-D float32 array.

Uses decord’s AudioReader (already a project dependency for video decoding) to produce a mono waveform at the requested sampling rate.

nemo_rl.data.datasets.utils.assert_no_double_bos(
token_ids: torch.Tensor,
tokenizer: nemo_rl.data.datasets.utils.TokenizerType,
) → None#

Assert that there are no double starting BOS tokens in the message.

Parameters:

  • token_ids – List of token IDs

  • tokenizer – Tokenizer

nemo_rl.data.datasets.utils.pil_to_base64(image: PIL.Image.Image, format: str = 'PNG') → str#

Converts a PIL Image object to a base64 encoded string.

Parameters:

  • image – The PIL Image object to convert.

  • format – The image format (e.g., “PNG”, “JPEG”). Defaults to “PNG”.

Returns:

A base64 encoded string representation of the image.

nemo_rl.data.datasets.utils.load_dataset_from_path(
data_path: str,
data_subset: Optional[str] = None,
data_split: Optional[str] = 'train',
)#

Load a dataset from a local file, huggingface dataset, or Arrow dataset (saved with save_to_disk).

Parameters:

  • data_path – The path to the dataset.

  • data_subset – The subset to load from the dataset. Only supported for huggingface datasets.

  • data_split – The split to load from the dataset.

nemo_rl.data.datasets.utils.resolve_external_dataset_class(dataset_name: str) → type#

Resolve a fully-qualified dotted dataset path to a class.

Used by both load_response_dataset and load_preference_dataset to support user-defined datasets that live outside nemo_rl so users do not have to edit the built-in DATASET_REGISTRY to plug in their own dataset class. The class must be importable from PYTHONPATH (or the active virtual environment).

The caller is expected to have already verified that dataset_name looks like a dotted import path (i.e. contains a .); this helper focuses on the import / attribute-lookup / type-validation steps and raises ValueError with an actionable message on any failure.

nemo_rl.data.datasets.utils.update_single_dataset_config(
data_config: dict,
default_data_config: dict,
) → None#

Fill the single dataset config with default dataset config.

nemo_rl.data.datasets.utils.merge_datasets(datasets: list[Any]) → Any#

Merge map-style datasets while supporting non-HF wrappers.

Hugging Face’s concatenate_datasets only accepts HF Dataset objects. Some response datasets, such as PreservingDataset, intentionally bypass the HF schema machinery to preserve heterogeneous nested structures. When those datasets are present, fall back to a generic concatenation wrapper that still provides __len__ and integer __getitem__ for downstream processing.

nemo_rl.data.datasets.utils.extract_necessary_env_names(data_config: dict) → list[str]#

Extract the necessary environment names from the data config.

Some environments are set in env_configs but not used in the data config. This function extracts the necessary environment names from the data config.

Parameters:

data_config – The data config.

Returns:

The necessary environment names.

nemo_rl.data.datasets.utils.get_huggingface_cache_path(
repo_id,
branch='main',
repo_type='datasets',
)#