nemo_rl.data.interfaces#

Module Contents#

Classes#

DatumSpec

DPODatumSpec

TaskDataSpec

TaskDataProcessFnCallable

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

Data#

LLMMessageLogType

FlatMessagesType

PathLike

TokenizerType

API#

nemo_rl.data.interfaces.LLMMessageLogType#

None

nemo_rl.data.interfaces.FlatMessagesType#

None

nemo_rl.data.interfaces.PathLike#

None

nemo_rl.data.interfaces.TokenizerType#

None

class nemo_rl.data.interfaces.DatumSpec[source]#

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: NotRequired[str]#

None

stop_strings: NotRequired[list[str]]#

None

__extra__: NotRequired[Any]#

None

class nemo_rl.data.interfaces.DPODatumSpec[source]#

Bases: typing.TypedDict

message_log_chosen: nemo_rl.data.interfaces.LLMMessageLogType#

None

message_log_rejected: nemo_rl.data.interfaces.LLMMessageLogType#

None

length_chosen: int#

None

length_rejected: int#

None

loss_multiplier: float#

None

idx: int#

None

class nemo_rl.data.interfaces.TaskDataSpec[source]#
task_name: Optional[str]#

None

prompt_file: Optional[nemo_rl.data.interfaces.PathLike]#

None

system_prompt_file: Optional[nemo_rl.data.interfaces.PathLike]#

None

__post_init__() → None[source]#
copy_defaults(from_spec: nemo_rl.data.interfaces.TaskDataSpec) → None[source]#

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

class nemo_rl.data.interfaces.TaskDataProcessFnCallable[source]#

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: nemo_rl.data.interfaces.TokenizerType,
max_seq_length: int | None,
idx: int,
) → nemo_rl.data.interfaces.DatumSpec[source]#