`nemo_rl.data.energon.multimodal.task_encoders.base`#

Module Contents#

Classes#

BaseSFTTaskEncoder

Common SFT lifecycle shared by the Energon task encoders.

Data#

`SFTCooker`
`__all__`

API#

nemo_rl.data.energon.multimodal.task_encoders.base.SFTCooker: TypeAlias#: None

class nemo_rl.data.energon.multimodal.task_encoders.base.BaseSFTTaskEncoder( *, cooker_functions: collections.abc.Sequence[nemo_rl.data.energon.multimodal.task_encoders.base.SFTCooker], )#

Bases: megatron.energon.DefaultTaskEncoder[nemo_rl.data.energon.multimodal.types.CanonicalSFTSample, nemo_rl.data.energon.multimodal.types.EncodedSFTSample, nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any], nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any]], abc.ABC

Common SFT lifecycle shared by the Energon task encoders.

Initialization

abstractmethod preencode_sample( sample: nemo_rl.data.energon.multimodal.types.CanonicalSFTSample, ) → nemo_rl.data.energon.multimodal.types.EncodedSFTSample#: Encode one sample before optional pack selection.

abstractmethod postencode_sample( sample: nemo_rl.data.energon.multimodal.types.EncodedSFTSample, ) → nemo_rl.data.energon.multimodal.types.EncodedSFTSample#: Finish the selected sample before physical packing.

abstractmethod batch( samples: list[nemo_rl.data.energon.multimodal.types.EncodedSFTSample], ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#: Combine encoded samples into one minibatch.

abstractmethod encode_batch( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any], ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#: Finish one minibatch before the loader emits it.

nemo_rl.data.energon.multimodal.task_encoders.base.__all__#: [‘BaseSFTTaskEncoder’, ‘SFTCooker’]