"""Nemotron-Dense Audex audio-understanding model for HuggingFace inference.""" from __future__ import annotations from typing import Optional import torch from torch import nn import torch.nn.functional as F from transformers.cache_utils import DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from .configuration_nemotron_h_audio import NemotronDenseAudexConfig from .modeling_nemotron_dense import NemotronDenseForCausalLM class NemotronDenseAudexRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-5): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.float() variance = hidden_states.pow(2).mean(dim=-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.eps) return (self.weight.float() * hidden_states).to(input_dtype) class NemotronDenseAudexProjector(nn.Module): """Megatron sound_projection equivalent for TP1 HF inference.""" def __init__(self, config: NemotronDenseAudexConfig): super().__init__() self.norm = NemotronDenseAudexRMSNorm( config.audio_encoder_hidden_size, eps=config.audio_projector_norm_eps, ) self.fc1 = nn.Linear( config.audio_encoder_hidden_size, config.audio_projector_intermediate_size, bias=False, ) self.fc2 = nn.Linear( config.audio_projector_intermediate_size, config.hidden_size, bias=False, ) self.activation = config.audio_projector_activation def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.norm(hidden_states) hidden_states = self.fc1(hidden_states) if self.activation == "relu2": hidden_states = F.relu(hidden_states).pow(2) elif self.activation == "gelu": hidden_states = F.gelu(hidden_states) else: raise ValueError(f"Unsupported audio projector activation: {self.activation}") return self.fc2(hidden_states) def _build_qwen2_audio_encoder(audio_config: dict) -> nn.Module: try: from transformers.models.qwen2_audio.configuration_qwen2_audio import Qwen2AudioEncoderConfig from transformers.models.qwen2_audio.modeling_qwen2_audio import Qwen2AudioEncoder except Exception as exc: # pragma: no cover - version/environment guard raise ImportError( "Qwen2-Audio support is required for NV-Whisper. " "Install a transformers build that provides transformers.models.qwen2_audio." ) from exc cfg = Qwen2AudioEncoderConfig(**dict(audio_config)) return Qwen2AudioEncoder(cfg) class NemotronDenseAudexForConditionalGeneration(NemotronDenseForCausalLM): """Nemotron-Dense CausalLM plus NV-Whisper encoder and sound projection. State dict layout intentionally keeps the baseline LLM key names: `model.*` and `lm_head.*` load exactly as in the LLM-only checkpoint. New audio tensors live under `audio_encoder.*` and `audio_projector.*`. """ config_class = NemotronDenseAudexConfig _tp_plan = None _base_model_tp_plan = None base_model_tp_plan = None def __init__(self, config: NemotronDenseAudexConfig): super().__init__(config) self.audio_encoder = _build_qwen2_audio_encoder(config.audio_config) self.audio_projector = NemotronDenseAudexProjector(config) def encode_audio(self, input_features: torch.Tensor) -> torch.Tensor: """Encode Whisper features into LLM hidden-space audio embeddings.""" encoder_param = next(self.audio_encoder.parameters()) input_features = input_features.to(device=encoder_param.device, dtype=encoder_param.dtype) encoder_outputs = self.audio_encoder(input_features=input_features, return_dict=True) audio_hidden = encoder_outputs.last_hidden_state projector_param = next(self.audio_projector.parameters()) audio_hidden = audio_hidden.to(device=projector_param.device, dtype=projector_param.dtype) return self.audio_projector(audio_hidden) def _audio_embeddings_by_sample( self, input_features: Optional[torch.Tensor], audio_embeddings: Optional[torch.Tensor], batch_size: int, ) -> list[torch.Tensor]: if audio_embeddings is None: if input_features is None: raise ValueError("input_features or audio_embeddings must be provided for audio injection") if input_features.ndim == 3: projected = self.encode_audio(input_features) if batch_size != 1: raise ValueError( "3D input_features represent a single sample. " "Use 4D (batch, clips, mel_bins, frames) features for batched audio." ) return [projected.reshape(-1, projected.shape[-1])] if input_features.ndim == 4: bsz, clips, mel_bins, frames = input_features.shape if bsz != batch_size: raise ValueError(f"input_features batch {bsz} != input_ids batch {batch_size}") flat_features = input_features.reshape(bsz * clips, mel_bins, frames) projected = self.encode_audio(flat_features) projected = projected.reshape(bsz, clips * projected.shape[1], projected.shape[-1]) return [projected[idx] for idx in range(bsz)] raise ValueError(f"Expected 3D or 4D input_features, got {tuple(input_features.shape)}") if audio_embeddings.ndim == 2: if batch_size != 1: raise ValueError("2D audio_embeddings only support batch_size=1") return [audio_embeddings.to(device=self.device)] if audio_embeddings.ndim == 3: if audio_embeddings.shape[0] != batch_size: raise ValueError(f"audio_embeddings batch {audio_embeddings.shape[0]} != input_ids batch {batch_size}") return [audio_embeddings[idx].to(device=self.device) for idx in range(batch_size)] raise ValueError(f"Expected 2D or 3D audio_embeddings, got {tuple(audio_embeddings.shape)}") def prepare_inputs_embeds( self, input_ids: torch.LongTensor, input_features: Optional[torch.Tensor] = None, audio_embeddings: Optional[torch.Tensor] = None, ) -> torch.Tensor: if input_ids is None: raise ValueError("input_ids are required when injecting audio embeddings") if self.config.sound_token_id is None: raise ValueError("config.sound_token_id is required for audio embedding injection") embed_device = self.model.embed_tokens.weight.device input_ids = input_ids.to(embed_device) inputs_embeds = self.model.embed_tokens(input_ids).clone() audio_by_sample = self._audio_embeddings_by_sample( input_features=input_features, audio_embeddings=audio_embeddings, batch_size=input_ids.shape[0], ) for batch_idx, audio in enumerate(audio_by_sample): mask = input_ids[batch_idx].to(self.device) == self.config.sound_token_id mask = mask.to(embed_device) expected = int(mask.sum().item()) if expected != audio.shape[0]: raise ValueError( "Mismatch between token count and projected audio tokens: " f"sample={batch_idx} placeholders={expected} audio_tokens={audio.shape[0]}" ) inputs_embeds[batch_idx, mask] = audio.to(device=embed_device, dtype=inputs_embeds.dtype) return inputs_embeds @staticmethod def _is_prefill(past_key_values) -> bool: if past_key_values is None: return True get_len = getattr(past_key_values, "get_seq_length", None) if callable(get_len): return get_len() == 0 return len(past_key_values) == 0 def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, input_features=None, audio_embeddings=None, cache_position=None, position_ids=None, use_cache=True, **kwargs, ): if self._is_prefill(past_key_values) and inputs_embeds is None and ( input_features is not None or audio_embeddings is not None ): inputs_embeds = self.prepare_inputs_embeds( input_ids=input_ids, input_features=input_features, audio_embeddings=audio_embeddings, ) return super().prepare_inputs_for_generation( input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, cache_position=cache_position, position_ids=position_ids, use_cache=use_cache, **kwargs, ) def _dense_forward_from_embeds( self, inputs_embeds: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, use_cache: Optional[bool] = None, return_dict: Optional[bool] = None, ): use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if use_cache and past_key_values is None: past_key_values = DynamicCache() hidden_states = inputs_embeds for decoder_layer in self.model.layers: hidden_states = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, ) hidden_states = self.model.norm(hidden_states) if not return_dict: return tuple(v for v in [hidden_states, past_key_values] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) def forward( self, input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, input_features: Optional[torch.Tensor] = None, audio_embeddings: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ): if inputs_embeds is None and self._is_prefill(past_key_values) and ( input_features is not None or audio_embeddings is not None ): inputs_embeds = self.prepare_inputs_embeds( input_ids=input_ids, input_features=input_features, audio_embeddings=audio_embeddings, ) input_ids = None if inputs_embeds is None: return super().forward( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, labels=labels, use_cache=use_cache, return_dict=return_dict, **kwargs, ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self._dense_forward_from_embeds( inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, return_dict=return_dict, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) loss = None if labels is not None: labels = labels.to(logits.device) shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1)) if not return_dict: output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, )