nemo_rl.data.interfaces

Module Contents

Classes

DatumSpec
TaskDataSpec
TaskDataProcessFnCallable	A callable that processes a loaded datum dictionary into a DatumSpec.

Data

LLMMessageLogType
FlatMessagesType

API

nemo_rl.data.interfaces.LLMMessageLogType

None

nemo_rl.data.interfaces.FlatMessagesType

None

class nemo_rl.data.interfaces.DatumSpec

Bases: typing.TypedDict

message_log: nemo_rl.data.interfaces.LLMMessageLogType

None

length: int

None

extra_env_info: Dict[str, Any]

None

loss_multiplier: float

None

idx: int

None

task_name: Optional[str]

‘default’

stop_strings: Optional[List[str]]

None

__extra__: Any

None

class nemo_rl.data.interfaces.TaskDataSpec

task_name: Optional[str]

None

prompt_file: Optional[os.PathLike]

None

system_prompt_file: Optional[Union[str, os.PathLike]]

None

__post_init__()

copy_defaults(from_spec: nemo_rl.data.interfaces.TaskDataSpec)

Apply default values from another Task instance for any None attributes.

class nemo_rl.data.interfaces.TaskDataProcessFnCallable

Bases: typing.Protocol

A callable that processes a loaded datum dictionary into a DatumSpec.

abstractmethod __call__(

datum_dict: Dict[str, Any],

task_data_spec: nemo_rl.data.interfaces.TaskDataSpec,

tokenizer,

max_seq_length: int,

idx: int,

) → nemo_rl.data.interfaces.DatumSpec