`bridge.models.nemotron_omni.data.collate_fn`#

Shared Nemotron Omni collation for Direct-HF and Energon datasets.

Module Contents#

Functions#

`_build_padded_assistant_loss_masks`	Build assistant loss masks without treating batch padding as message boundaries.
`_validate_nemotron_omni_visual_keys`	Validate the model-owned visual input contract retained for API compatibility.
`_validate_reserved_media_text`	Reject raw text that collides with processor-owned media placeholders.
`_pad_text_rows`	Right-pad unbatched token rows without deriving padding from token values.
`_pil_images`	Normalize one image/frame payload to a flat list of PIL images.
`_tensor_image_to_uint8`	Convert one CHW image tensor in either [0, 1] or [0, 255] range to uint8.
`_decode_video_path`
`_video_frames`	Decode raw/path video payloads or flatten already-decoded frame payloads.
`_patchify_frame`	Apply the public normalization kernel on MCore’s square temporal canvas.
`_pixel_shuffled_token_count`	Return RADIO tokens after the model’s fixed 2x2 spatial pixel shuffle.
`_join_content_parts`	Separate structured parts without duplicating source-owned newlines.
`_render_text_conversation`	Replace structured image parts with literal placeholders in source order.
`_prepare_temporal_rows`	Build row-local temporal prompts and one packed all-frame vision tensor.
`_prepare_standard_rows`	Use the HF processor for text/images while preserving row ownership.
`_audio_waveform`
`_add_audio_inputs`	Extract audio features and align each row’s sound placeholder count.
`_adjust_image_placeholders`
`_owned_media_positions`	Validate processor-owned media groups and return their placeholder positions.
`_split_media_token_counts`	Split flattened per-media token counts back into their structured rows.
`_build_media_token_validity_mask`	Validate reserved placeholders and mark processor-owned media anchors.
`_pack_dynamic_images`
`_model_merge_row_lengths`	Return per-row lengths after MCore replaces compact image placeholders.
`_nonpacked_multimodal_compact_width`	Choose a compact width whose model-side merged width stays in bounds.
`_pack_omni_rows_to_mcore_thd`	Pack compact Omni rows using their lengths after model-side embedding merge.
`nemotron_omni_collate_fn`	Build one model-ready Omni batch from either HF or Energon examples.
`nemotron_omni_llava_collate_fn`	Collate inputs for the deprecated LLaVA collapse/expand path.
`nemotron_omni_expanded_collate_fn`	Collate processor-expanded inputs for the canonical Nemotron Omni model.

Data#

`CHATML_ASSISTANT_START`
`CHATML_ASSISTANT_END`
`CHATML_OTHER_ROLE_STARTS`
`VISION_FRAME_SIZE`
`PIXEL_SHUFFLE_FACTOR`
`_NEMOTRON_OMNI_VISUAL_KEYS`
`_RESERVED_MEDIA_TEXT`

API#

bridge.models.nemotron_omni.data.collate_fn.CHATML_ASSISTANT_START#: ‘<|im_start|>assistant\n’

bridge.models.nemotron_omni.data.collate_fn.CHATML_ASSISTANT_END#: ‘<|im_end|>\n’

bridge.models.nemotron_omni.data.collate_fn.CHATML_OTHER_ROLE_STARTS#: None

bridge.models.nemotron_omni.data.collate_fn.VISION_FRAME_SIZE#: 512

bridge.models.nemotron_omni.data.collate_fn.PIXEL_SHUFFLE_FACTOR#: 2

bridge.models.nemotron_omni.data.collate_fn._NEMOTRON_OMNI_VISUAL_KEYS#: (‘pixel_values’,)

bridge.models.nemotron_omni.data.collate_fn._RESERVED_MEDIA_TEXT#: None

bridge.models.nemotron_omni.data.collate_fn._build_padded_assistant_loss_masks( examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], input_ids: torch.Tensor, attention_mask: torch.Tensor, processor: Any, skipped_tokens: torch.Tensor, *, boundary_config: megatron.bridge.data.conversation_processing.AssistantMaskBoundaryConfig, ) → torch.Tensor#: Build assistant loss masks without treating batch padding as message boundaries.

bridge.models.nemotron_omni.data.collate_fn._validate_nemotron_omni_visual_keys( visual_keys: object = None, ) → None#: Validate the model-owned visual input contract retained for API compatibility.

bridge.models.nemotron_omni.data.collate_fn._validate_reserved_media_text( examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], ) → None#: Reject raw text that collides with processor-owned media placeholders.

bridge.models.nemotron_omni.data.collate_fn._pad_text_rows( rows: collections.abc.Sequence[torch.Tensor], *, pad_token_id: int, ) → tuple[torch.Tensor, torch.Tensor]#: Right-pad unbatched token rows without deriving padding from token values.

bridge.models.nemotron_omni.data.collate_fn._pil_images(payload: Any) → list[Any]#: Normalize one image/frame payload to a flat list of PIL images.

bridge.models.nemotron_omni.data.collate_fn._tensor_image_to_uint8(image: torch.Tensor) → numpy.ndarray#: Convert one CHW image tensor in either [0, 1] or [0, 255] range to uint8.

bridge.models.nemotron_omni.data.collate_fn._decode_video_path( path: str, *, video_fps: float, video_nframes: int, ) → tuple[list[Any], transformers.video_utils.VideoMetadata]#

bridge.models.nemotron_omni.data.collate_fn._video_frames( payload: Any, *, video_fps: float, video_nframes: int, ) → tuple[list[Any], transformers.video_utils.VideoMetadata]#: Decode raw/path video payloads or flatten already-decoded frame payloads.

bridge.models.nemotron_omni.data.collate_fn._patchify_frame( frame: Any, *, height: int, width: int, patch_dim: int, ) → torch.Tensor#

Apply the public normalization kernel on MCore’s square temporal canvas.

The public HF processor preserves video aspect ratio, but pinned MCore’s temporal path currently stacks tubelets and pixel-shuffles them as a common square grid. Until MCore supports ragged non-square tubelets, Bridge keeps the required 512-square compatibility canvas while matching HF’s antialiased bicubic interpolation and RADIO normalization.

bridge.models.nemotron_omni.data.collate_fn._pixel_shuffled_token_count( *, height: int, width: int, patch_dim: int, ) → int#: Return RADIO tokens after the model’s fixed 2x2 spatial pixel shuffle.

bridge.models.nemotron_omni.data.collate_fn._join_content_parts(parts: collections.abc.Sequence[str]) → str#: Separate structured parts without duplicating source-owned newlines.

bridge.models.nemotron_omni.data.collate_fn._render_text_conversation( example: collections.abc.Mapping[str, Any], ) → tuple[list[dict[str, Any]], list[Any]]#: Replace structured image parts with literal placeholders in source order.

bridge.models.nemotron_omni.data.collate_fn._prepare_temporal_rows( examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], processor: Any, *, temporal_patch_size: int, video_fps: float, video_nframes: int, patch_dim: int, temporal_video_resize_mode: Literal[fixed_512, bridge.models.nemotron_omni.data.collate_fn._prepare_temporal_rows.processor], ) → tuple[dict[str, Any], list[dict[str, Any]], torch.Tensor, list[int]]#: Build row-local temporal prompts and one packed all-frame vision tensor.

bridge.models.nemotron_omni.data.collate_fn._prepare_standard_rows( examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], processor: Any, ) → tuple[dict[str, Any], list[dict[str, Any]], torch.Tensor | None, list[int]]#: Use the HF processor for text/images while preserving row ownership.

bridge.models.nemotron_omni.data.collate_fn._audio_waveform( example: collections.abc.Mapping[str, Any], *, target_sampling_rate: int = 16000, ) → numpy.ndarray | None#

bridge.models.nemotron_omni.data.collate_fn._add_audio_inputs( batch: dict[str, Any], examples: collections.abc.Sequence[collections.abc.Mapping[str, Any]], processor: Any, *, max_audio_duration: float, num_mel_bins: int, ) → list[list[int]]#: Extract audio features and align each row’s sound placeholder count.

bridge.models.nemotron_omni.data.collate_fn._adjust_image_placeholders( batch: dict[str, Any], loss_mask: torch.Tensor, processor: Any, num_tiles: torch.Tensor | None, ) → tuple[dict[str, torch.Tensor], torch.Tensor]#

bridge.models.nemotron_omni.data.collate_fn._owned_media_positions( token_row: torch.Tensor, *, start_token_id: int, end_token_id: int, placeholder_token_id: int, expected_token_counts: collections.abc.Sequence[int], modality: str, ) → torch.Tensor#: Validate processor-owned media groups and return their placeholder positions.

bridge.models.nemotron_omni.data.collate_fn._split_media_token_counts( token_counts: torch.Tensor | collections.abc.Sequence[int], groups_per_row: collections.abc.Sequence[int], *, modality: str, ) → list[list[int]]#: Split flattened per-media token counts back into their structured rows.

bridge.models.nemotron_omni.data.collate_fn._build_media_token_validity_mask( batch: collections.abc.Mapping[str, Any], processor: Any, *, image_token_counts: collections.abc.Sequence[collections.abc.Sequence[int]], sound_token_counts: collections.abc.Sequence[collections.abc.Sequence[int]], ) → torch.Tensor#: Validate reserved placeholders and mark processor-owned media anchors.

bridge.models.nemotron_omni.data.collate_fn._pack_dynamic_images( batch: dict[str, Any], *, patch_dim: int, ) → None#

bridge.models.nemotron_omni.data.collate_fn._model_merge_row_lengths( batch: collections.abc.Mapping[str, Any], processor: Any, *, use_per_image_token_counts: bool, patch_dim: int, ) → torch.Tensor#: Return per-row lengths after MCore replaces compact image placeholders.

bridge.models.nemotron_omni.data.collate_fn._nonpacked_multimodal_compact_width( batch: collections.abc.Mapping[str, Any], post_merge_row_lengths: torch.Tensor, *, sequence_length: int, pad_to_max_length: bool, pad_to_multiple_of: int, ) → int#: Choose a compact width whose model-side merged width stays in bounds.

bridge.models.nemotron_omni.data.collate_fn._pack_omni_rows_to_mcore_thd( batch: dict[str, Any], post_merge_row_lengths: torch.Tensor, *, sequence_length: int | None, pad_to_max_length: bool, pad_to_multiple_of: int, pad_token_id: int, ) → None#: Pack compact Omni rows using their lengths after model-side embedding merge.

bridge.models.nemotron_omni.data.collate_fn.nemotron_omni_collate_fn( examples: list[collections.abc.Mapping[str, Any]], processor: Any, start_of_response_token: Any = None, *, visual_keys: object = None, min_pixels: int | None = None, max_pixels: int | None = None, enable_in_batch_packing: bool = False, sequence_length: int | None = None, pad_to_max_length: bool = False, pad_to_multiple_of: int = 128, in_batch_packing_pad_to_multiple_of: int = 1, max_audio_duration: float = 30.0, num_mel_bins: int = 128, temporal_patch_size: int = 2, video_fps: float = 1.0, video_nframes: int = 8, use_temporal_video_embedder: bool = False, patch_dim: int = 16, temporal_video_resize_mode: Literal[fixed_512, bridge.models.nemotron_omni.data.collate_fn.nemotron_omni_collate_fn.processor] = 'fixed_512', collapse_image_tokens: bool = False, ) → dict[str, Any]#

Build one model-ready Omni batch from either HF or Energon examples.

The canonical :class:NemotronOmniModel consumes the processor-expanded token sequence, with one image placeholder for every projected feature. temporal_video_resize_mode="processor" uses the image processor’s aspect-preserving video grid and derives each tubelet’s placeholder count from its returned size metadata. Direct collator calls retain the "fixed_512" default; the Energon task encoder opts into processor mode by default. Video frame labels use source-frame timestamps and do not add an implicit prose prefix. Use :func:nemotron_omni_llava_collate_fn for the legacy LLaVA collapse/expand contract.

bridge.models.nemotron_omni.data.collate_fn.nemotron_omni_llava_collate_fn(

*args,

**kwargs,

) → dict[str, torch.Tensor]#: Collate inputs for the deprecated LLaVA collapse/expand path.

bridge.models.nemotron_omni.data.collate_fn.nemotron_omni_expanded_collate_fn(

*args,

**kwargs,

) → dict[str, Any]#: Collate processor-expanded inputs for the canonical Nemotron Omni model.

bridge.models.nemotron_omni.data.collate_fn#

Module Contents#

Functions#

Data#

API#

`bridge.models.nemotron_omni.data.collate_fn`#