`nemo_rl.data`#

Subpackages#

Submodules#

Package Contents#

Classes#

`DataConfig`
`MathDataConfig`

API#

class nemo_rl.data.DataConfig[source]#

Bases: typing.TypedDict

max_input_seq_length: int#: None

prompt_file: NotRequired[str]#: None

system_prompt_file: NotRequired[str]#: None

dataset_name: str#: None

val_dataset_name: NotRequired[str]#: None

add_bos: NotRequired[bool]#: None

add_eos: NotRequired[bool]#: None

input_key: NotRequired[str]#: None

output_key: NotRequired[str]#: None

add_generation_prompt: NotRequired[bool]#: None

add_system_prompt: NotRequired[bool]#: None

split: NotRequired[str]#: None

shuffle: NotRequired[bool]#: None

seed: NotRequired[int]#: None

download_dir: NotRequired[str]#: None

train_data_path: NotRequired[str]#: None

val_data_paths: NotRequired[dict[str, str]]#: None

num_workers: NotRequired[int]#: None

class nemo_rl.data.MathDataConfig[source]#

Bases: nemo_rl.data.DataConfig

problem_key: str#: None

solution_key: str#: None