`nemo_rl.data.energon.multimodal.packing`#

Energon-owned selection and materialization of multimodal SFT packs.

Module Contents#

`_cost`
`select_samples_to_pack`	Group compatible sources and run the configured packer.
`pack_selected_samples`	Turn one selected source group into a physical pack.
`prepare_packed_sft_batch`	Create model tensors for a batch of physical Energon packs.

nemo_rl.data.energon.multimodal.packing._cost( sample: nemo_rl.data.energon.multimodal.types.EncodedSFTSample, multiple: int, ) → int#

nemo_rl.data.energon.multimodal.packing.select_samples_to_pack( samples: list[nemo_rl.data.energon.multimodal.types.EncodedSFTSample], *, packer: nemo_rl.data.packing.SequencePacker, sequence_length_pad_multiple: int, ) → list[list[nemo_rl.data.energon.multimodal.types.EncodedSFTSample]]#: Group compatible sources and run the configured packer.

nemo_rl.data.energon.multimodal.packing.pack_selected_samples( samples: list[nemo_rl.data.energon.multimodal.types.EncodedSFTSample], *, pack_capacity: int, sequence_length_pad_multiple: int, ) → nemo_rl.data.energon.multimodal.types.PackedSFTSample#: Turn one selected source group into a physical pack.

nemo_rl.data.energon.multimodal.packing.prepare_packed_sft_batch( packs: list[nemo_rl.data.energon.multimodal.types.PackedSFTSample], *, tokenizer: transformers.PreTrainedTokenizerBase, only_unmask_final: bool, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#: Create model tensors for a batch of physical Energon packs.

nemo_rl.data.energon.multimodal.packing.__all__#: [‘pack_selected_samples’, ‘prepare_packed_sft_batch’, ‘select_samples_to_pack’]