`bridge.data.energon.nemotron_omni_task_encoder`#

Nemotron Omni Energon adapter backed by the shared HF-style collator.

Module Contents#

Classes#

`_LegacyVisualTensorMapping`	Live mapping view over a batch’s `GenericVisualInputs` fields.
`NemotronOmniTaskBatch`	HF-style batch with the legacy `visual_tensors` constructor/property.
`NemotronOmniTaskEncoder`	Normalize Energon samples and delegate all model processing to one collator.

Data#

`NemotronOmniTaskSample`
`_VISUAL_INPUT_FIELDS`

API#

bridge.data.energon.nemotron_omni_task_encoder.NemotronOmniTaskSample#: None

bridge.data.energon.nemotron_omni_task_encoder._VISUAL_INPUT_FIELDS#: ‘frozenset(…)’

class bridge.data.energon.nemotron_omni_task_encoder._LegacyVisualTensorMapping( batch: bridge.data.energon.nemotron_omni_task_encoder.NemotronOmniTaskBatch, )#

Bases: collections.abc.MutableMapping[str, typing.Any]

Live mapping view over a batch’s GenericVisualInputs fields.

Initialization

__getitem__(key: str) → Any#

__setitem__(key: str, value: Any) → None#

__delitem__(key: str) → None#

__iter__() → collections.abc.Iterator[str]#

__len__() → int#

class bridge.data.energon.nemotron_omni_task_encoder.NemotronOmniTaskBatch(

*args: Any,

visual_inputs: megatron.bridge.training.utils.visual_inputs.GenericVisualInputs | None = None,

visual_tensors: Mapping[str, Any] | None = None,

**kwargs: Any,

)#

Bases: megatron.bridge.data.energon.hf_task_encoder.HFEnergonBatch

HF-style batch with the legacy visual_tensors constructor/property.

Initialization

num_patches: torch.Tensor | None#: None

sound_clips: torch.Tensor | None#: None

sound_length: torch.Tensor | None#: None

imgs_sizes: torch.Tensor | None#: None

num_frames: torch.Tensor | None#: None

num_image_tiles: torch.Tensor | None#: None

property visual_tensors: collections.abc.MutableMapping[str, Any]#: Expose a live legacy visual tensor mapping for low-level batch consumers.

class bridge.data.energon.nemotron_omni_task_encoder.NemotronOmniTaskEncoder( processor: Any, seq_length: int = 4096, max_audio_duration: float = 30.0, num_mel_bins: int = 128, visual_keys: Sequence[str] = ('pixel_values',), temporal_patch_size: int = 2, video_fps: float = 1.0, video_nframes: int = 8, use_temporal_video_embedder: bool = False, patch_dim: int = 16, pad_to_max_length: bool = False, pad_to_multiple_of: int = 128, enable_in_batch_packing: bool = False, in_batch_packing_pad_to_multiple_of: int = 1, collapse_image_tokens: bool = False, )#

Bases: megatron.bridge.data.energon.hf_task_encoder.HFTaskEncoder

Normalize Energon samples and delegate all model processing to one collator.

The task encoder owns only source adaptation and configuration. Tokenization, assistant masking, modality-token expansion, padding, and in-batch packing are performed by the canonical expanded-sequence collator for both Direct-HF and Energon datasets. collapse_image_tokens=True selects the deprecated LLaVA compatibility contract.

Initialization

collate_fn( examples: list[dict[str, Any]], ) → dict[str, Any]#: Collate normalized Energon examples with the shared Omni path.

batch( samples: list[megatron.bridge.data.energon.hf_task_encoder.HFEnergonSample], ) → bridge.data.energon.nemotron_omni_task_encoder.NemotronOmniTaskBatch#: Collate shared samples while retaining the model-specific batch type.

encode_batch( batch: megatron.bridge.data.energon.hf_task_encoder.HFEnergonBatch, ) → dict[str, Any]#: Return the shared batch plus the legacy tokens alias.

bridge.data.energon.nemotron_omni_task_encoder#

Module Contents#

Classes#

Data#

API#

`bridge.data.energon.nemotron_omni_task_encoder`#