nemo_rl.data.hf_datasets.refcoco#

Module Contents#

Classes#

RefCOCODataset

Functions#

download_and_unzip

Downloads a zip file from a given URL to a target directory and unzips it into a specified subdirectory within the target directory, showing download progress.

format_refcoco_dataset

Format the RefCOCO dataset from huggingface.

prepare_refcoco_dataset

API#

nemo_rl.data.hf_datasets.refcoco.download_and_unzip(
url: str,
target_directory: str,
subdir_name: str = '.',
)#

Downloads a zip file from a given URL to a target directory and unzips it into a specified subdirectory within the target directory, showing download progress.

Parameters:

  • url (str) – The URL of the zip file to download.

  • target_directory (str) – The directory where the zip file will be downloaded and unzipped.

  • subdir_name (str) – The name of the subdirectory within the target_directory where the contents of the zip file will be unzipped. Defaults to “train”.

nemo_rl.data.hf_datasets.refcoco.format_refcoco_dataset(
example: dict[str, Any],
width: int = 256,
height: int = 256,
caption_type: str = 'random',
prompt_file: Optional[str] = None,
) → dict[str, Any]#

Format the RefCOCO dataset from huggingface.

This should be replaced with our own curated RefCOCO/+/g dataset soon

Parameters:

  • example – The example to format.

  • width – The width of the resized image.

  • height – The height of the resized image.

  • caption_type – The type of caption to use.

nemo_rl.data.hf_datasets.refcoco.prepare_refcoco_dataset(
split: str = 'default',
task_name: Optional[str] = None,
path_to_coco_images: Optional[Union[str, pathlib.Path]] = None,
)#
class nemo_rl.data.hf_datasets.refcoco.RefCOCODataset(
split: str = 'default',
prompt_file: Optional[str] = None,
download_dir: Optional[str] = None,
)#

Initialization

Simple wrapper around the RefCOCO dataset.

Parameters:

  • split – The split of the dataset to use (currently only ‘default’ is supported)

  • prompt_file – The file containing the prompt for the dataset.