diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..52373fe24473b1aa44333d318f578ae6bf04b49b 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
diff --git a/audio_model.py b/audio_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..98dd17a7c74a22749c86c04d36858450b2dc53c7
--- /dev/null
+++ b/audio_model.py
@@ -0,0 +1,174 @@
+# 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.
+
+"""Sound/Audio model components for multimodal integration.
+
+This module provides the SoundEncoder (wrapping Parakeet from HuggingFace transformers)
+and SoundProjection (MLP to project audio embeddings to LLM hidden size).
+
+The Parakeet model in HuggingFace transformers is documented at:
+https://huggingface.co/docs/transformers/en/model_doc/parakeet
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from transformers import ParakeetEncoder, ParakeetEncoderConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class SquaredReLU(nn.Module):
+ """Squared ReLU activation function."""
+ def forward(self, x):
+ return torch.pow(torch.nn.functional.relu(x), 2)
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-5):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.eps = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+ return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+
+class SoundProjection(nn.Module):
+ """MLP projection from sound encoder hidden size to LLM hidden size.
+
+ Architecture: RMSNorm -> linear1 -> SquaredReLU -> linear2
+
+ This matches the Megatron checkpoint conversion structure:
+ - sound_projection.norm.weight
+ - sound_projection.linear1.weight
+ - sound_projection.linear2.weight
+ - sound_projection.linear1.bias (optional)
+ - sound_projection.linear2.bias (optional)
+ """
+
+ def __init__(
+ self,
+ sound_hidden_size: int,
+ projection_hidden_size: int,
+ llm_hidden_size: int,
+ bias: bool = True,
+ eps: float = 1e-5,
+ ):
+ super().__init__()
+ self.norm = RMSNorm(sound_hidden_size, eps=eps)
+ self.linear1 = nn.Linear(sound_hidden_size, projection_hidden_size, bias=bias)
+ self.activation = SquaredReLU()
+ self.linear2 = nn.Linear(projection_hidden_size, llm_hidden_size, bias=bias)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ """Project sound embeddings to LLM embedding space.
+
+ Args:
+ hidden_states: Sound encoder output [batch, seq_len, sound_hidden_size]
+
+ Returns:
+ Projected embeddings [batch, seq_len, llm_hidden_size]
+ """
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.linear1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.linear2(hidden_states)
+ return hidden_states
+
+
+class SoundEncoder(nn.Module):
+ """Wrapper around the Parakeet encoder from HuggingFace transformers.
+
+ The Parakeet model is an ASR model with a Fast Conformer encoder.
+ We use only the encoder portion to extract audio embeddings.
+
+ Checkpoint structure:
+ - sound_encoder.encoder.feature_extractor.* -> Feature extraction (mel spectrogram)
+ - sound_encoder.encoder.pre_encode.* -> Pre-encoding convolutions
+ - sound_encoder.encoder.layers.* -> Conformer layers
+
+ Reference: https://huggingface.co/docs/transformers/en/model_doc/parakeet
+ """
+
+ def __init__(self, config=None):
+ super().__init__()
+
+ if config is not None:
+ # Build from config - handle both dict and config object
+ if hasattr(config, '__dict__'):
+ # It's a config object, extract relevant params for ParakeetConfig
+ config_dict = {
+ 'attention_bias': getattr(config, 'attention_bias', False),
+ 'hidden_size': getattr(config, 'hidden_size', 1024),
+ 'num_attention_heads': getattr(config, 'num_attention_heads', 8),
+ 'num_hidden_layers': getattr(config, 'num_hidden_layers', 24),
+ 'intermediate_size': getattr(config, 'intermediate_size', 4096),
+ 'conv_kernel_size': getattr(config, 'conv_kernel_size', 31),
+ 'convolution_bias': getattr(config, 'convolution_bias', False),
+ 'feat_in': getattr(config, 'feat_in', 80),
+ 'subsampling_factor': getattr(config, 'subsampling_factor', 8),
+ 'subsampling_conv_channels': getattr(config, 'subsampling_conv_channels', 256),
+ 'subsampling_conv_kernel_size': getattr(config, 'subsampling_conv_kernel_size', 3),
+ 'subsampling_conv_stride': getattr(config, 'subsampling_conv_stride', 2),
+ 'num_mel_bins': getattr(config, 'num_mel_bins', 128),
+ 'scale_input': getattr(config, 'scale_input', False),
+ }
+ elif isinstance(config, dict):
+ config_dict = config
+ else:
+ config_dict = {}
+
+ # Create ParakeetConfig with the extracted parameters
+ parakeet_config = ParakeetEncoderConfig(**config_dict)
+ self.config = parakeet_config
+ self.encoder = ParakeetEncoder(parakeet_config)
+ else:
+ raise ValueError(
+ "config must be provided, "
+ "and ParakeetEncoder must be available in transformers."
+ )
+
+ def forward(
+ self,
+ input_features: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Encode audio features.
+
+ Args:
+ input_features: Mel spectrogram features [batch, seq_len, feature_dim]
+ attention_mask: Optional attention mask [batch, seq_len]
+
+ Returns:
+ Audio embeddings [batch, encoded_seq_len, hidden_size]
+ """
+ outputs = self.encoder(
+ input_features=input_features,
+ attention_mask=attention_mask,
+ )
+ # Return the last hidden state
+ return outputs.last_hidden_state
+
+ @property
+ def hidden_size(self) -> int:
+ """Return the hidden size of the encoder."""
+ return self.config.hidden_size
diff --git a/chat_template.jinja b/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..6381c9df849b06a13c04a9992ed9cfe3f1e25a81
--- /dev/null
+++ b/chat_template.jinja
@@ -0,0 +1,273 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+ {%- if json_dict is mapping %}
+ {%- for json_key in json_dict if json_key not in handled_keys %}
+ {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+ {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '' }}
+ {%- else %}
+ {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '' }}
+ {%- endif %}
+ {%- endfor %}
+ {%- endif %}
+{%- endmacro -%}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set reasoning_budget = reasoning_budget if reasoning_budget is defined else None %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{#- Scan messages for VLM thinking toggles to override enable_thinking -#}
+{%- set toggle = namespace(enable=enable_thinking) %}
+{%- for m in messages %}
+ {%- if m['role'] == 'user' or m['role'] == 'system' -%}
+ {%- if m['content'] is string -%}
+ {%- set c = m['content'] %}
+ {%- if '/think' in c.replace('', '') -%}
+ {%- set toggle.enable = true -%}
+ {%- elif '/no_think' in c -%}
+ {%- set toggle.enable = false -%}
+ {%- endif -%}
+ {%- else -%}
+ {%- for part in m['content'] -%}
+ {%- if part['type'] == 'text' -%}
+ {%- set c = part['text'] %}
+ {%- if '/think' in c.replace('', '') -%}
+ {%- set toggle.enable = true -%}
+ {%- elif '/no_think' in c -%}
+ {%- set toggle.enable = false -%}
+ {%- endif -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- endif -%}
+ {%- endif -%}
+{%- endfor -%}
+{#- Prepare message iteration similar to LM template -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+ {%- if m["role"] == "user" %}
+ {%- set ns.last_user_idx = loop.index0 %}
+ {%- endif %}
+{%- endfor -%}
+
+{%- if messages[0]["role"] == "system" %}
+ {%- set system_message = messages[0]["content"] %}
+ {%- set loop_messages = messages[1:] %}
+{%- else %}
+ {%- set system_message = "" %}
+ {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+ {%- set tools = [] %}
+{%- endif %}
+{#- Recompute last_user_idx relative to loop_messages after handling system -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+ {%- if m["role"] == "user" %}
+ {%- set ns.last_user_idx = loop.index0 %}
+ {%- endif %}
+{%- endfor -%}
+{#- System preamble with LM formatting, sanitize thinking toggles -#}
+{%- if system_message is defined %}
+ {%- set sys_content = system_message | string %}
+ {%- set sys_content = sys_content.replace('', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '') %}
+ {{- "<|im_start|>system\n" + sys_content }}
+{%- else %}
+ {%- if tools is iterable and tools | length > 0 %}
+ {{- "<|im_start|>system\n" }}
+ {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+ {%- if system_message is defined and system_message | length > 0 %}
+ {{- "\n\n" }}
+ {%- endif %}
+ {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
+ {{- "" }}
+ {%- for tool in tools %}
+ {%- if tool.function is defined %}
+ {%- set tool = tool.function %}
+ {%- endif %}
+ {{- "\n\n" ~ tool.name ~ "" }}
+ {%- if tool.description is defined %}
+ {{- '\n' ~ (tool.description | trim) ~ '' }}
+ {%- endif %}
+ {{- '\n' }}
+ {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+ {%- for param_name, param_fields in tool.parameters.properties|items %}
+ {{- '\n' }}
+ {{- '\n' ~ param_name ~ '' }}
+ {%- if param_fields.type is defined %}
+ {{- '\n' ~ (param_fields.type | string) ~ '' }}
+ {%- endif %}
+ {%- if param_fields.description is defined %}
+ {{- '\n' ~ (param_fields.description | trim) ~ '' }}
+ {%- endif %}
+ {%- if param_fields.enum is defined %}
+ {{- '\n' ~ (param_fields.enum | tojson | safe) ~ '' }}
+ {%- endif %}
+ {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+ {{- render_extra_keys(param_fields, handled_keys) }}
+ {{- '\n' }}
+ {%- endfor %}
+ {%- endif %}
+ {%- set handled_keys = ['type', 'properties', 'required'] %}
+ {{- render_extra_keys(tool.parameters, handled_keys) }}
+ {%- if tool.parameters is defined and tool.parameters.required is defined %}
+ {{- '\n' ~ (tool.parameters.required | tojson | safe) ~ '' }}
+ {%- endif %}
+ {{- '\n' }}
+ {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+ {{- render_extra_keys(tool, handled_keys) }}
+ {{- '\n' }}
+ {%- endfor %}
+ {{- "\n" }}
+
+ {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }}
+{%- endif -%}
+{%- if system_message is defined %}
+ {{- '<|im_end|>\n' }}
+{%- else %}
+ {%- if tools is iterable and tools | length > 0 %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+{%- endif -%}
+{#- Iterate conversation -#}
+{%- for message in loop_messages %}
+ {%- if message.role == "assistant" %}
+ {#- Use LM assistant handling -#}
+ {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+ {%- set content = "\n" ~ message.reasoning_content ~ "\n\n" ~ (message.content | default('', true)) %}
+ {%- else %}
+ {%- set content = message.content | default('', true) %}
+ {%- if content is string -%}
+ {%- if '' not in content and '' not in content -%}
+ {%- set content = "" ~ content -%}
+ {%- endif -%}
+ {%- else -%}
+ {%- set content = content -%}
+ {%- endif -%}
+ {%- endif %}
+ {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+ {{- '<|im_start|>assistant\n' }}
+ {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+ {%- if content is string and content | trim | length > 0 %}
+ {%- if include_content %}
+ {{- (content | trim) ~ '\n' -}}
+ {%- else %}
+ {%- set c = (content | string) %}
+ {%- if '' in c %}
+ {%- set c = c.split('')[-1] %}
+ {%- elif '' in c %}
+ {%- set c = c.split('')[0] %}
+ {%- endif %}
+ {%- set c = "" ~ c | trim %}
+ {%- if c | length > 0 %}
+ {{- c ~ '\n' -}}
+ {%- endif %}
+ {%- endif %}
+ {%- else %}
+ {{- "" -}}
+ {%- endif %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if tool_call.function is defined %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n\n' -}}
+ {%- if tool_call.arguments is defined %}
+ {%- for args_name, args_value in tool_call.arguments|items %}
+ {{- '\n' -}}
+ {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+ {{- args_value ~ '\n\n' -}}
+ {%- endfor %}
+ {%- endif %}
+ {{- '\n\n' -}}
+ {%- endfor %}
+ {{- '<|im_end|>\n' }}
+ {%- else %}
+ {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+ {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+ {%- else %}
+ {%- set c = (content | default('', true) | string) %}
+ {%- if '' in c and '' in c %}
+ {%- set c = "" ~ c.split('')[-1] %}
+ {%- endif %}
+ {%- set c = c | trim %}
+ {%- if c | length > 0 %}
+ {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+ {%- else %}
+ {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+ {%- endif %}
+ {%- elif message.role == "user" or message.role == "system" %}
+ {{- '<|im_start|>' + message.role + '\n' }}
+ {#- Build VLM multimodal content when content is a sequence -#}
+ {%- if message.content is string -%}
+ {%- set content = (message.content | string) %}
+ {%- else -%}
+ {%- set text_ns = namespace(val='') -%}
+ {%- set mm_content = '' -%}
+ {%- set counters = namespace(images=0, videos=0, audios=0) -%}
+ {%- for part in message['content'] -%}
+ {%- if part['type'] == 'image' or part['type'] == 'image_url' -%}
+ {%- set counters.images = counters.images + 1 -%}
+ {%- elif part['type'] == 'video' or part['type'] == 'video_url' -%}
+ {%- set counters.videos = counters.videos + 1 -%}
+ {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' -%}
+ {%- set counters.audios = counters.audios + 1 -%}
+ {%- elif part['type'] == 'text' -%}
+ {%- set text_ns.val = text_ns.val + part['text'] -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- if '' in text_ns.val -%}
+ {%- set counters.images = 0 -%}
+ {%- endif -%}
+ {%- if '', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '') -%}
+ {%- set content = content | trim -%}
+ {%- if message.role == "user" and loop.index0 == ns.last_user_idx and reasoning_budget is not none -%}
+ {{- content + '\n\n{thinking token budget: ' + (reasoning_budget | string) + '}' -}}
+ {%- else -%}
+ {{- content -}}
+ {%- endif -%}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
+ {{- '<|im_start|>user\n' }}
+ {%- endif %}
+ {{- '\n' }}
+ {{- message.content }}
+ {{- '\n\n' }}
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
+ {{- '<|im_end|>\n' }}
+ {%- elif loop.last %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- else %}
+ {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+ {%- endif %}
+{%- endfor -%}
+{#- Generation prompt using computed thinking toggle -#}
+{%- if add_generation_prompt %}
+ {%- if toggle.enable %}
+ {{- '<|im_start|>assistant\n\n' }}
+ {%- else %}
+ {{- '<|im_start|>assistant\n' }}
+ {%- endif %}
+{%- endif %}
\ No newline at end of file
diff --git a/config.json b/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..8854c4c9b8115b85ae96d1769ca9b27ac181c82d
--- /dev/null
+++ b/config.json
@@ -0,0 +1,30394 @@
+{
+ "architectures": [
+ "NemotronH_Nano_Omni_Reasoning_V3"
+ ],
+ "auto_map": {
+ "AutoConfig": "configuration.NemotronH_Nano_Omni_Reasoning_V3_Config",
+ "AutoModel": "modeling.NemotronH_Nano_Omni_Reasoning_V3",
+ "AutoModelForCausalLM": "modeling.NemotronH_Nano_Omni_Reasoning_V3"
+ },
+ "max_sequence_length": 131072,
+ "downsample_ratio": 0.5,
+ "force_image_size": 512,
+ "patch_size": 16,
+ "use_thumbnail": true,
+ "eos_token_id": 11,
+ "model_type": "NemotronH_Nano_Omni_Reasoning_V3",
+ "ps_version": "v2",
+ "template": "n5h_5p5_nanov2",
+ "torch_dtype": "bfloat16",
+ "image_tag_type": "internvl",
+ "img_context_token_id": 18,
+ "video_context_token_id": 131081,
+ "img_context_token": "",
+ "video_context_token": "