checkpoint_folder_full/modeling_nemotron_h_audio.py

"""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 <so_embedding> 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,
        )