# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/nemotron3_5_asr/modular_nemotron3_5_asr.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_nemotron3_5_asr.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Inc. team. 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. import math from dataclasses import dataclass from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMode from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available, logging from ..auto import AutoModel from .configuration_nemotron3_5_asr import Nemotron3_5AsrConfig from .generation_nemotron3_5_asr import Nemotron3_5AsrGenerationMixin, Nemotron3_5AsrRNNTDecoderCache if is_torch_available(): import torch from torch import nn logger = logging.get_logger(__name__) @dataclass class Nemotron3_5AsrRNNTOutput(BaseModelOutputWithPooling): """ encoder_past_key_values (`Cache`, *optional*): Updated encoder attention K/V sliding-window cache, returned when encoding audio with `use_cache=True` (cache-aware streaming). Pass it to the next chunk's forward. padding_cache (`NemotronAsrStreamingEncoderCausalConvPaddingCache`, *optional*): Updated unified streaming conv cache (subsampling Conv2d + conformer depthwise Conv1d), returned when encoding audio with `use_cache=True`. Pass it to the next chunk's forward. """ loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None decoder_cache: Nemotron3_5AsrRNNTDecoderCache | None = None encoder_past_key_values: Cache | None = None padding_cache: "NemotronAsrStreamingEncoderCausalConvPaddingCache | None" = None # noqa: F821 @auto_docstring class Nemotron3_5AsrPreTrainedModel(PreTrainedModel): config: Nemotron3_5AsrConfig base_model_prefix = "model" main_input_name = "input_features" input_modalities = "audio" supports_gradient_checkpointing = True _no_split_modules = None _supports_flat_attention_mask = True _supports_sdpa = True # flex attention is incompatible as this model uses a float attention mask (relative position bias) across the board _supports_flex_attn = False # TODO: @eustlb, add support when flash attention supports custom attention bias _supports_flash_attn = False _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = {} @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) def _get_subsampling_output_length(self, input_lengths: torch.Tensor): encoder_config = getattr(self.config, "encoder_config", self.config) kernel_size = encoder_config.subsampling_conv_kernel_size stride = encoder_config.subsampling_conv_stride num_layers = int(math.log2(encoder_config.subsampling_factor)) # The subsampling Conv2d is always causal: NeMo's CausalConv2D pads (left=kernel-1, right=stride-1). all_paddings = (kernel_size - 1) + (stride - 1) add_pad = all_paddings - kernel_size lengths = input_lengths for _ in range(num_layers): lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + 1.0 lengths = torch.floor(lengths) return lengths.to(dtype=torch.int) def _get_output_attention_mask(self, attention_mask: torch.Tensor, target_length: int | None = None): """ Convert the input attention mask to its subsampled form. `target_length` sets the desired output length, useful when the attention mask length differs from `sum(-1).max()` (i.e., when the longest sequence in the batch is padded) """ output_lengths = self._get_subsampling_output_length(attention_mask.sum(-1)) # Use target_length if provided, otherwise use max length in batch max_length = target_length if target_length is not None else output_lengths.max() attention_mask = torch.arange(max_length, device=attention_mask.device) < output_lengths[:, None] return attention_mask class Nemotron3_5AsrPromptProjector(nn.Module): def __init__(self, config: Nemotron3_5AsrConfig): super().__init__() self.linear_1 = nn.Linear( config.encoder_config.hidden_size + config.num_prompts, config.prompt_intermediate_size ) self.act = nn.ReLU() self.linear_2 = nn.Linear(config.prompt_intermediate_size, config.encoder_config.hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.linear_1(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.linear_2(hidden_states) return hidden_states class Nemotron3_5AsrRNNTDecoder(nn.Module): """LSTM-based prediction network For RNN-T""" def __init__(self, config: Nemotron3_5AsrConfig): super().__init__() self.blank_token_id = config.blank_token_id self.embedding = nn.Embedding(config.vocab_size, config.decoder_hidden_size) self.lstm = nn.LSTM( input_size=config.decoder_hidden_size, hidden_size=config.decoder_hidden_size, num_layers=config.num_decoder_layers, batch_first=True, ) self.decoder_projector = nn.Linear(config.decoder_hidden_size, config.decoder_hidden_size) def forward( self, input_ids: torch.LongTensor, cache: Nemotron3_5AsrRNNTDecoderCache | None = None, ) -> torch.Tensor: if cache is not None: blank_mask = input_ids[:, -1] == self.blank_token_id # All-blank fast path: skip decoder when all batch elements predict blank if cache.is_initialized and blank_mask.all(): return cache.cache embeddings = self.embedding(input_ids) # Get cached hidden/cell states if available, otherwise initialize with Nemotron3_5AsrRNNTDecoderCache if cache is not None: was_initialized = cache.is_initialized if not was_initialized: cache.lazy_initialization(embeddings) hidden_cell_states = (cache.hidden_state, cache.cell_state) else: hidden_cell_states = None lstm_output, (hidden_state, cell_state) = self.lstm(embeddings, hidden_cell_states) decoder_output = self.decoder_projector(lstm_output) if cache is not None: mask = ~blank_mask if was_initialized else None cache.update(decoder_output, hidden_state, cell_state, mask=mask) return cache.cache return decoder_output class Nemotron3_5AsrRNNTJointNetwork(nn.Module): """Joint network that combines encoder and decoder outputs to predict token logits.""" def __init__(self, config: Nemotron3_5AsrConfig): super().__init__() self.activation = ACT2FN[config.hidden_act] self.head = nn.Linear(config.decoder_hidden_size, config.vocab_size) self.vocab_size = config.vocab_size def forward( self, decoder_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: joint_output = self.activation(encoder_hidden_states + decoder_hidden_states) return self.head(joint_output) @auto_docstring( custom_intro=""" Nemotron3_5Asr Encoder with an RNN-T (Recurrent Neural Network Transducer) head and language-ID prompt conditioning. """ ) class Nemotron3_5AsrForRNNT(Nemotron3_5AsrPreTrainedModel, Nemotron3_5AsrGenerationMixin): config: Nemotron3_5AsrConfig _no_split_modules = ["Nemotron3_5AsrRNNTDecoder"] _supported_generation_modes = [GenerationMode.GREEDY_SEARCH] def __init__(self, config: Nemotron3_5AsrConfig): super().__init__(config) self.encoder = AutoModel.from_config(config.encoder_config) self.encoder_projector = nn.Linear(config.encoder_config.hidden_size, config.decoder_hidden_size) self.decoder = Nemotron3_5AsrRNNTDecoder(config) self.joint = Nemotron3_5AsrRNNTJointNetwork(config) self.max_symbols_per_step = config.max_symbols_per_step # used in generation self.prompt_projector = Nemotron3_5AsrPromptProjector(config) self.post_init() @can_return_tuple def get_audio_features( self, input_features: torch.Tensor, attention_mask: torch.Tensor | None = None, prompt_ids: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: encoder_outputs = self.encoder( input_features=input_features, attention_mask=attention_mask, **kwargs, ) hidden_states = encoder_outputs.last_hidden_state if prompt_ids is None: logger.warning_once( "`prompt_ids` not provided; defaulting to " f"`config.default_prompt_id={self.config.default_prompt_id}` (auto language detection). " "Pass `language` to the processor or `prompt_ids` to condition on a specific language." ) prompt_ids = torch.full( (hidden_states.shape[0],), self.config.default_prompt_id, dtype=torch.long, device=hidden_states.device, ) prompt_ids = prompt_ids.to(hidden_states.device) one_hot = nn.functional.one_hot(prompt_ids, num_classes=self.config.num_prompts).to(hidden_states.dtype) one_hot = one_hot[:, None, :].expand(-1, hidden_states.shape[1], -1) fused = self.prompt_projector(torch.cat([hidden_states, one_hot], dim=-1)) encoder_outputs.pooler_output = self.encoder_projector(fused) return encoder_outputs @auto_docstring @can_return_tuple def forward( self, input_features: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, decoder_input_ids: torch.LongTensor | None = None, decoder_cache: Nemotron3_5AsrRNNTDecoderCache | None = None, use_decoder_cache: bool | None = None, encoder_outputs: BaseModelOutputWithPooling | None = None, labels: torch.Tensor | None = None, num_lookahead_tokens: int | None = None, prompt_ids: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> Nemotron3_5AsrRNNTOutput: r""" decoder_input_ids (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*): Decoder input token ids for single-step inference. decoder_cache (`Nemotron3_5AsrRNNTDecoderCache`, *optional*): Decoder LSTM cache. Reused on blank predictions to skip the LSTM step. use_decoder_cache (`bool`, *optional*): Whether to allocate and use a decoder cache when none is provided. encoder_outputs (`tuple(torch.FloatTensor)`, *optional*): Pre-computed encoder outputs (last_hidden_state, pooler_output, ...). num_lookahead_tokens (`int`, *optional*): Right attention context (lookahead, in subsampled encoder frames) forwarded to the encoder. Defaults to `config.encoder_config.default_num_lookahead_tokens`. prompt_ids (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Language-prompt indices for language-ID conditioning. Produced by the processor from `language`. Turned into the broadcast one-hot consumed by `prompt_projector`. Example: ```python >>> from transformers import AutoProcessor, Nemotron3_5AsrForRNNT >>> from datasets import load_dataset, Audio >>> model_id = "nvidia/nemotron-3.5-asr-streaming-0.6b" >>> processor = AutoProcessor.from_pretrained(model_id) >>> model = Nemotron3_5AsrForRNNT.from_pretrained(model_id) >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate)) >>> inputs = processor(ds[0]["audio"]["array"], language="en-US") >>> outputs = model(**inputs) ``` """ if encoder_outputs is None: encoder_outputs = self.get_audio_features( input_features=input_features, attention_mask=attention_mask, num_lookahead_tokens=num_lookahead_tokens, prompt_ids=prompt_ids, **kwargs, ) if use_decoder_cache and decoder_cache is None: decoder_cache = Nemotron3_5AsrRNNTDecoderCache() decoder_hidden_states = self.decoder(decoder_input_ids, cache=decoder_cache) logits = self.joint( encoder_hidden_states=encoder_outputs.pooler_output[:, :, None, :], decoder_hidden_states=decoder_hidden_states[:, None, :, :], ).squeeze(2) loss = None if labels is not None: loss = self.loss_function(logits=logits, labels=labels, encoder_outputs=encoder_outputs) return Nemotron3_5AsrRNNTOutput( loss=loss, logits=logits, last_hidden_state=encoder_outputs.last_hidden_state, pooler_output=encoder_outputs.pooler_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, decoder_cache=decoder_cache, encoder_past_key_values=encoder_outputs.past_key_values, padding_cache=encoder_outputs.padding_cache, ) __all__ = ["Nemotron3_5AsrRNNTOutput", "Nemotron3_5AsrForRNNT", "Nemotron3_5AsrPreTrainedModel"]