File size: 7,140 Bytes
63b7dc2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from huggingface_hub.dataclasses import strict
from transformers.configuration_utils import PreTrainedConfig
from transformers.models.nemotron_h import NemotronHConfig
from transformers.models.parakeet.configuration_parakeet import ParakeetEncoderConfig
from transformers.utils import auto_docstring, logging
from .configuration_radio import RadioConfig
__all__ = ["NemotronH_Omni_Reasoning_V3_Config"]
logger = logging.get_logger(__name__)
def _nemotron_h_compatible_config(config: dict) -> dict:
"""Translate current Nemotron-H layer names for Transformers versions using legacy names."""
config = dict(config)
legacy_names = {"linear_attention": "mamba", "full_attention": "attention"}
for field in ("layers_block_type", "mtp_layers_block_type"):
if field in config and config[field] is not None:
config[field] = [legacy_names.get(layer_type, layer_type) for layer_type in config[field]]
return config
@auto_docstring(checkpoint="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16")
@strict
class NemotronH_Omni_Reasoning_V3_Config(PreTrainedConfig):
r"""
vision_config (`dict` or `RadioConfig`, *optional*):
Configuration for the RADIO vision encoder. Defaults to a default [`RadioConfig`].
llm_config (`dict` or `NemotronHConfig`, *optional*):
Configuration for the NemotronH language model. Defaults to a default [`NemotronHConfig`].
sound_config (`dict` or `ParakeetEncoderConfig`, *optional*):
Configuration for the optional Parakeet sound encoder. `None` disables the audio branch.
force_image_size (`int`, *optional*):
Fixed input image resolution (in pixels) the vision tower expects.
downsample_ratio (`float`, *optional*, defaults to 0.5):
Pixel-shuffle spatial downsample ratio applied to the vision features.
template (`str`, *optional*):
Conversation template name (legacy; the chat template is carried by the processor).
ps_version (`str`, *optional*, defaults to `"v1"`):
Pixel-shuffle implementation version (`"v2"` swaps height/width back; `"v1"` is the legacy
transposed variant).
image_tag_type (`str`, *optional*, defaults to `"internvl"`):
Image-tag convention used when expanding image placeholders.
projector_hidden_size (`int`, *optional*, defaults to 4096):
Hidden size of the vision-to-LLM MLP projector.
vit_hidden_size (`int`, *optional*, defaults to 1280):
Hidden size of the RADIO vision features.
video_pruning_rate (`float`, *optional*, defaults to 0.0):
Efficient-Video-Sampling token pruning rate; `0.0` disables pruning.
video_temporal_patch_size (`int`, *optional*, defaults to 2):
Number of frames collapsed into a single temporal patch by the video embedder.
patch_size (`int`, *optional*, defaults to 16):
Vision patch size in pixels.
img_context_token_id (`int`, *optional*):
Token id used as the image-context placeholder in `input_ids`.
img_context_token (`str`, *optional*, defaults to `"<image>"`):
Textual image-context placeholder token.
video_context_token_id (`int`, *optional*):
Token id used as the video-context placeholder in `input_ids`.
video_context_token (`str`, *optional*, defaults to `"<video>"`):
Textual video-context placeholder token.
sound_context_token_id (`int`, *optional*):
Token id used as the audio-context placeholder in `input_ids`.
sound_context_token (`str`, *optional*, defaults to `"<audio>"`):
Textual audio-context placeholder token.
"""
model_type = "nemotron_h_omni"
is_composition = True
vision_config: dict | RadioConfig | None = None
llm_config: dict | NemotronHConfig | None = None
sound_config: dict | ParakeetEncoderConfig | None = None
force_image_size: int | None = None
downsample_ratio: float = 0.5
template: str | None = None # TODO move out of here and into the tokenizer
ps_version: str = "v1" # Pixel shuffle version
image_tag_type: str = "internvl" # TODO: into the tokenizer too?
projector_hidden_size: int = 4096
vit_hidden_size: int = 1280
video_pruning_rate: float = 0.0
video_temporal_patch_size: int = 2
# Vision token settings
patch_size: int = 16
img_context_token_id: int | None = None
img_context_token: str = "<image>"
video_context_token_id: int | None = None
video_context_token: str = "<video>"
# Sound/audio token settings
sound_context_token_id: int | None = None
sound_context_token: str = "<audio>"
def __post_init__(self, **kwargs):
if isinstance(self.vision_config, dict):
self.vision_config = RadioConfig(**self.vision_config)
elif self.vision_config is None:
self.vision_config = RadioConfig()
# The vision tower needs the temporal patch size to build its video patch projection.
self.vision_config.video_temporal_patch_size = self.video_temporal_patch_size
# Handle both cases: when loading from JSON (llm_config is dict) and when called
# internally by transformers (llm_config is None).
if isinstance(self.llm_config, dict):
self.llm_config = NemotronHConfig(**_nemotron_h_compatible_config(self.llm_config))
elif self.llm_config is None:
self.llm_config = NemotronHConfig()
# This checkpoint's sound_config omits `attention_bias`/`scale_input`, which are `False` for
# its Parakeet variant, so supply them before building the encoder config.
if isinstance(self.sound_config, dict):
sound_config = {"attention_bias": False, "scale_input": False, **self.sound_config}
sound_config.pop("model_type", None)
self.sound_config = ParakeetEncoderConfig(**sound_config)
super().__post_init__(**kwargs)
# `attn_implementation` flows in through `**kwargs` (the base `PreTrainedConfig` stores it as
# `self._attn_implementation`, as for every other model); propagate it to the language model.
self.llm_config._attn_implementation = self._attn_implementation
# vLLM's `NemotronH_Nano_VL_V2` implementation reads the language-model sub-config as
# `config.text_config`. Our HF config stores it as `config.llm_config`; expose an alias so the
# same config object loads under both loaders without having to duplicate the dict on disk.
@property
def text_config(self):
return self.llm_config
|