nemo_rl.data.hf_datasets.geometry3k#

Module Contents#

Classes#

Geometry3KDataset

Functions#

format_geometry3k_dataset

Format the Geometry3K dataset into an OpenAI-API-like message log.

prepare_geometry3k_dataset

API#

nemo_rl.data.hf_datasets.geometry3k.format_geometry3k_dataset(
example: dict[str, Any],
return_pil: bool = False,
) → dict[str, Any]#

Format the Geometry3K dataset into an OpenAI-API-like message log.

nemo_rl.data.hf_datasets.geometry3k.prepare_geometry3k_dataset(
split: str = 'train',
task_name: str = 'geometry3k',
)#
class nemo_rl.data.hf_datasets.geometry3k.Geometry3KDataset(
split: str = 'train',
prompt_file: Optional[str] = None,
)#

Initialization

Simple wrapper around the Geometry3K dataset.

Parameters:

  • split – The split of the dataset to use.

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