#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/granite_speech5/modular_granite_speech5.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_granite_speech5.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 IBM and The HuggingFace 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.

from collections.abc import Callable
from dataclasses import dataclass

import torch
from torch import nn

from ... import initialization as init
from ...activations import ACT2FN
from ...generation import CompileConfig, GenerationMixin
from ...modeling_layers import GradientCheckpointingLayer
from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutput
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from ...processing_utils import Unpack
from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple
from ...utils.generic import merge_with_config_defaults
from ...utils.output_capturing import capture_outputs
from ..auto import AutoModel
from .configuration_granite_speech5 import GraniteSpeech5CTCConfig, GraniteSpeech5EncoderConfig


@auto_docstring(
    custom_intro="""
    Extends [~modeling_outputs.BaseModelOutputWithPooling] to include the output attention mask since sequence length
    is not preserved in the model's forward.
    """
)
@dataclass
class GraniteSpeech5EncoderModelOutput(BaseModelOutputWithPooling):
    r"""
    attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
        Mask to avoid performing attention on padding token indices after sequence compression. Returned because the
        sequence length may differ from the input sequence length. Mask values selected in `[0, 1]`:

        - 1 for tokens that are **not masked**,
        - 0 for tokens that are **masked**.
    """

    attention_mask: torch.Tensor | None = None


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


class GraniteSpeech5EncoderAttention(nn.Module):
    """Block-wise self-attention with Shaw's relative positional embeddings."""

    def __init__(self, config: GraniteSpeech5EncoderConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = False
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True)
        self.context_size = config.context_size
        self.rel_pos_emb = nn.Embedding(2 * config.max_position_embeddings + 1, config.head_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, None]:
        batch_size, seq_length, _ = hidden_states.shape

        # right-pad the sequence to a whole number of blocks
        num_padded = -seq_length % self.context_size
        if num_padded > 0:
            hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, num_padded))
            if attention_mask is None:
                attention_mask = torch.ones(batch_size, seq_length, dtype=torch.bool, device=hidden_states.device)
            attention_mask = nn.functional.pad(attention_mask, (0, num_padded), value=False)

        # every block is folded into the batch dimension, so all blocks attend independently in one call
        flattened_batch_size = batch_size * (hidden_states.shape[1] // self.context_size)
        hidden_shape = (flattened_batch_size, self.context_size, self.config.num_attention_heads, -1)
        query_states = self.q_proj(hidden_states).reshape(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).reshape(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).reshape(hidden_shape).transpose(1, 2)

        relative_position_embeddings = self.rel_pos_emb(position_embeddings) * self.scaling
        queries = query_states.permute(2, 0, 1, 3).reshape(self.context_size, -1, self.head_dim)
        position_bias = queries @ relative_position_embeddings.transpose(1, 2)
        position_bias = position_bias.view(
            self.context_size, flattened_batch_size, self.config.num_attention_heads, self.context_size
        )
        position_bias = position_bias.permute(1, 2, 0, 3)
        position_bias = position_bias.contiguous()  # the fused attention kernels want a contiguous bias
        if attention_mask is not None:
            # mask padded key columns so no query attends to a padded frame
            key_mask = attention_mask.reshape(flattened_batch_size, self.context_size)
            position_bias = position_bias.masked_fill(
                ~key_mask[:, None, None, :], torch.finfo(position_bias.dtype).min
            )

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )
        attn_output, _ = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask=position_bias,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )
        attn_output = attn_output.reshape(batch_size, -1, self.config.num_attention_heads * self.head_dim)

        return self.o_proj(attn_output[:, :seq_length]), None


class GraniteSpeech5EncoderConvolutionModule(nn.Module):
    def __init__(self, config: GraniteSpeech5EncoderConfig, stride: int = 1):
        super().__init__()
        inner_dim = config.hidden_size * config.conv_expansion_factor

        self.pointwise_lin1 = nn.Linear(config.hidden_size, inner_dim * 2)
        self.depthwise_conv = nn.Conv1d(
            inner_dim,
            inner_dim,
            config.conv_kernel_size,  # kernel_size should be an odd number for 'SAME' padding
            stride=stride,
            padding=(config.conv_kernel_size - 1) // 2,
            groups=inner_dim,
            bias=False,
        )
        self.norm = nn.BatchNorm1d(inner_dim)
        self.pointwise_lin2 = nn.Linear(inner_dim, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        hidden_states = self.pointwise_lin1(hidden_states)
        hidden_states = nn.functional.glu(hidden_states, dim=-1)

        if attention_mask is not None:
            hidden_states = hidden_states.masked_fill(~attention_mask.unsqueeze(-1), 0.0)

        hidden_states = self.depthwise_conv(hidden_states.transpose(1, 2))
        hidden_states = nn.functional.silu(self.norm(hidden_states))
        hidden_states = self.pointwise_lin2(hidden_states.transpose(1, 2))
        return hidden_states


class GraniteSpeech5EncoderFeedForward(nn.Module):
    def __init__(self, config: GraniteSpeech5EncoderConfig):
        super().__init__()
        self.linear1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias)
        self.activation = ACT2FN[config.hidden_act]
        self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias)
        self.activation_dropout = config.activation_dropout

    def forward(self, hidden_states):
        hidden_states = self.activation(self.linear1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.linear2(hidden_states)
        return hidden_states


class GraniteSpeech5EncoderBlock(GradientCheckpointingLayer):
    def __init__(self, config: GraniteSpeech5EncoderConfig, layer_idx: int | None = None):
        super().__init__()
        self.gradient_checkpointing = False

        self.feed_forward1 = GraniteSpeech5EncoderFeedForward(config)
        self.self_attn = GraniteSpeech5EncoderAttention(config, layer_idx)
        self.conv = GraniteSpeech5EncoderConvolutionModule(config)
        self.feed_forward2 = GraniteSpeech5EncoderFeedForward(config)

        self.norm_feed_forward1 = nn.LayerNorm(config.hidden_size)
        self.norm_self_att = nn.LayerNorm(config.hidden_size)
        self.norm_conv = nn.LayerNorm(config.hidden_size)
        self.norm_feed_forward2 = nn.LayerNorm(config.hidden_size)
        self.norm_out = nn.LayerNorm(config.hidden_size)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states))
        hidden_states = residual + 0.5 * hidden_states  # the conformer architecture uses a factor of 0.5

        normalized_hidden_states = self.norm_self_att(hidden_states)
        attn_output, _ = self.self_attn(
            hidden_states=normalized_hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = hidden_states + attn_output

        conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask)
        hidden_states = hidden_states + conv_output

        ff2_output = self.feed_forward2(self.norm_feed_forward2(hidden_states))
        hidden_states = hidden_states + 0.5 * ff2_output  # the conformer architecture uses a factor of 0.5

        hidden_states = self.norm_out(hidden_states)

        return hidden_states


class GraniteSpeech5EncoderSubsamplingBlock(GraniteSpeech5EncoderBlock):
    """Conformer block that subsamples time by 2"""

    def __init__(self, config: GraniteSpeech5EncoderConfig, layer_idx: int | None = None):
        super().__init__(config, layer_idx)
        self.conv = GraniteSpeech5EncoderConvolutionModule(config, stride=2)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.feed_forward1(self.norm_feed_forward1(hidden_states))
        hidden_states = residual + 0.5 * hidden_states  # the conformer architecture uses a factor of 0.5

        normalized_hidden_states = self.norm_self_att(hidden_states)
        attn_output, _ = self.self_attn(
            hidden_states=normalized_hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = hidden_states + attn_output

        conv_output = self.conv(self.norm_conv(hidden_states), attention_mask=attention_mask)
        pooled = hidden_states.unfold(1, 2, 2).mean(-1)  # pool by 2 along time dim, eventually dropping last frame
        hidden_states = pooled + conv_output[:, : pooled.shape[1]]

        ff2_output = self.feed_forward2(self.norm_feed_forward2(hidden_states))
        hidden_states = hidden_states + 0.5 * ff2_output  # the conformer architecture uses a factor of 0.5

        hidden_states = self.norm_out(hidden_states)

        return hidden_states


@auto_docstring
class GraniteSpeech5PreTrainedModel(PreTrainedModel):
    config: GraniteSpeech5CTCConfig
    base_model_prefix = "model"
    main_input_name = "input_features"
    input_modalities = "audio"
    supports_gradient_checkpointing = True
    _no_split_modules = ["GraniteSpeech5EncoderBlock", "GraniteSpeech5EncoderSubsamplingBlock"]
    _supports_flat_attention_mask = True
    _supports_sdpa = True
    _supports_flex_attn = True

    # float attention bias is not supported by flash attention
    _supports_flash_attn = False

    _can_compile_fullgraph = True
    _supports_attention_backend = True

    _can_record_outputs = {
        "hidden_states": GraniteSpeech5EncoderBlock,
    }

    _keep_in_fp32_modules_strict = ["conv.norm"]

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, GraniteSpeech5Encoder):
            init.copy_(module.attention_dists, module.compute_attention_dists())

    def _get_subsampling_output_length(self, input_lengths: torch.Tensor):
        encoder_config = getattr(self.config, "encoder_config", self.config)
        output_lengths = input_lengths
        for _ in encoder_config.subsample_layers:
            output_lengths = output_lengths // 2
        return output_lengths

    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


def downsample_attention_mask(attention_mask: torch.Tensor) -> torch.Tensor:
    """Downsample a `(batch_size, seq_length)` padding mask by 2: a half-rate frame is valid iff both source
    frames are valid (a trailing odd frame is dropped, mirroring the residual pooling)."""
    half_length = attention_mask.shape[1] // 2
    return attention_mask[:, : 2 * half_length].reshape(attention_mask.shape[0], half_length, 2).all(dim=2)


@auto_docstring(
    custom_intro="""
    The Granite Speech 5.0 conformer encoder, adapted from the [Granite Speech CTC encoder](https://huggingface.co/papers/2505.08699)
    with block-wise time subsampling and self-conditioned CTC from the middle layer.
    """
)
class GraniteSpeech5Encoder(GraniteSpeech5PreTrainedModel):
    config: GraniteSpeech5EncoderConfig
    base_model_prefix = "encoder"

    def __init__(self, config: GraniteSpeech5EncoderConfig):
        super().__init__(config)
        self.gradient_checkpointing = False

        self.attention_dists = nn.Buffer(self.compute_attention_dists(), persistent=False)
        # see [`feature_extraction_granite_speech5.GraniteSpeech5FeatureExtractor`]
        # mel frames are delta-expanded (×2), then stacked in pairs (×2)
        self.input_linear = nn.Linear(config.num_mel_bins * 4, config.hidden_size, bias=True)
        self.layers = nn.ModuleList(
            [
                # CODEPATH: subsampling blocks
                GraniteSpeech5EncoderSubsamplingBlock(config, layer_idx)
                if layer_idx in config.subsample_layers
                else GraniteSpeech5EncoderBlock(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )

        self.out = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
        self.out_mid = nn.Linear(config.vocab_size, config.hidden_size, bias=True)

        self.post_init()

    @auto_docstring
    @merge_with_config_defaults
    @capture_outputs
    def forward(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        output_attention_mask: bool = True,
        **kwargs: Unpack[TransformersKwargs],
    ) -> GraniteSpeech5EncoderModelOutput:
        r"""
        output_attention_mask (`bool`, *optional*, defaults to `True`):
            Whether to return the output attention mask. Only effective when `attention_mask` is provided.

        Example:

        ```python
        >>> from transformers import AutoProcessor, GraniteSpeech5Encoder
        >>> from datasets import load_dataset, Audio

        >>> model_id = "ibm-granite/granite-speech-5.0-470m-turboctc"
        >>> processor = AutoProcessor.from_pretrained(model_id)
        >>> encoder = GraniteSpeech5Encoder.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"])
        >>> encoder_outputs = encoder(**inputs)

        >>> print(encoder_outputs.last_hidden_state.shape)
        ```
        """
        hidden_states = self.input_linear(input_features.to(self.input_linear.weight.dtype))
        if attention_mask is not None:
            attention_mask = attention_mask.bool()
            hidden_states = hidden_states.masked_fill(~attention_mask.unsqueeze(-1), 0.0)

        for layer_idx, layer in enumerate(self.layers):
            hidden_states = layer(
                hidden_states,
                attention_mask=attention_mask,
                position_embeddings=self.attention_dists,
                **kwargs,
            )

            # CODEPATH: the padding mask is halved after each subsampling block
            if layer_idx in self.config.subsample_layers and attention_mask is not None:
                attention_mask = downsample_attention_mask(attention_mask)

            # CODEPATH: self-conditioned CTC: feed the mid-layer CTC posteriors back into the hidden statess
            if layer_idx + 1 == self.config.num_hidden_layers // 2:
                mid_logits = self.out(hidden_states)
                mid_injection = self.out_mid(nn.functional.softmax(mid_logits, dim=-1))
                hidden_states = hidden_states + mid_injection.to(hidden_states.device)

        return GraniteSpeech5EncoderModelOutput(
            last_hidden_state=hidden_states,
            attention_mask=attention_mask.int() if attention_mask is not None and output_attention_mask else None,
        )

    def compute_attention_dists(self) -> torch.Tensor:
        """Clamped relative positional distances used by Shaw's relative positional embeddings."""
        context_size = self.config.context_size
        seq = torch.arange(context_size)
        relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
        return torch.clamp(relpos_dist, -context_size, context_size) + self.config.max_position_embeddings


@dataclass
class GraniteSpeech5CTCGenerateOutput(ModelOutput):
    """
    Outputs of GraniteSpeech5 CTC model generation.

    Args:
        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
            if all batches finished early due to the `eos_token_id`.
        logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
            Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
            at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
            each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
        attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
            `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
        hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
            `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
    """

    sequences: torch.LongTensor
    logits: tuple[torch.FloatTensor] | None = None
    attentions: tuple[tuple[torch.FloatTensor]] | None = None
    hidden_states: tuple[tuple[torch.FloatTensor]] | None = None


@auto_docstring(
    custom_intro="""
    Granite Speech 5.0 encoder with a Connectionist Temporal Classification (CTC) head.
    """
)
class GraniteSpeech5ForCTC(GraniteSpeech5PreTrainedModel, GenerationMixin):
    config: GraniteSpeech5CTCConfig
    # the CTC head is the same projection as the encoder's mid-layer self-conditioning head
    _tied_weights_keys = {
        "ctc_head.weight": "encoder.out.weight",
        "ctc_head.bias": "encoder.out.bias",
    }

    def __init__(self, config: GraniteSpeech5CTCConfig):
        super().__init__(config)
        self.encoder = AutoModel.from_config(config.encoder_config)
        self.ctc_head = nn.Linear(config.encoder_config.hidden_size, config.vocab_size, bias=True)

        self.post_init()

    @auto_docstring
    @can_return_tuple
    def forward(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> CausalLMOutput:
        r"""
        Example:

        ```python
        >>> from transformers import AutoProcessor, GraniteSpeech5ForCTC
        >>> from datasets import load_dataset, Audio

        >>> model_id = "nvidia/granite_speech5-ctc-1.1b"
        >>> processor = AutoProcessor.from_pretrained(model_id)
        >>> model = GraniteSpeech5ForCTC.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"], text=ds[0]["text"])
        >>> outputs = model(**inputs)

        >>> print(outputs.loss)
        ```"""

        if labels is not None:
            kwargs.setdefault("output_attention_mask", True)
        encoder_outputs = self.encoder(
            input_features=input_features,
            attention_mask=attention_mask,
            **kwargs,
        )

        hidden_states = encoder_outputs.last_hidden_state
        logits = self.ctc_head(hidden_states)

        loss = None
        if labels is not None:
            encoder_lengths = encoder_outputs.attention_mask.sum(-1)

            # assuming that padded tokens are filled with pad_token_id when not being attended to
            labels_mask = labels != self.config.pad_token_id
            target_lengths = labels_mask.sum(-1)
            flattened_targets = labels.masked_select(labels_mask)

            # ctc_loss doesn't support fp16
            log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)

            with torch.backends.cudnn.flags(enabled=False):
                loss = nn.functional.ctc_loss(
                    log_probs,
                    flattened_targets,
                    encoder_lengths,
                    target_lengths,
                    blank=self.config.pad_token_id,
                    reduction=self.config.ctc_loss_reduction,
                    zero_infinity=self.config.ctc_zero_infinity,
                )

        return CausalLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )

    @torch.no_grad()
    def generate(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        return_dict_in_generate: bool = False,
        compile_config: CompileConfig | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> GraniteSpeech5CTCGenerateOutput | torch.LongTensor:
        r"""
        compile_config ([`~generation.CompileConfig`], *optional*):
            If provided, `torch.compile` will be applied to the forward calls in the decoding loop.

        Example:

        ```python
        >>> from transformers import AutoProcessor, GraniteSpeech5ForCTC
        >>> from datasets import load_dataset, Audio

        >>> model_id = "ibm-granite/granite-speech-5.0-470m-turboctc"
        >>> processor = AutoProcessor.from_pretrained(model_id)
        >>> model = GraniteSpeech5ForCTC.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"])
        >>> predicted_ids = model.generate(**inputs)
        >>> transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)

        >>> print(transcription)
        ```
        """
        model_forward = self.get_compiled_call(compile_config) if compile_config is not None else self.__call__

        kwargs["return_dict"] = True
        outputs: CausalLMOutput = model_forward(
            input_features=input_features,
            attention_mask=attention_mask,
            **kwargs,
        )

        # greedy decoding
        sequences = outputs.logits.argmax(dim=-1)

        # mask out padded tokens
        if attention_mask is not None:
            attention_mask = self._get_output_attention_mask(attention_mask, target_length=sequences.shape[1])
            sequences[~attention_mask] = self.config.pad_token_id

        if return_dict_in_generate:
            return GraniteSpeech5CTCGenerateOutput(
                sequences=sequences,
                logits=outputs.logits,
                attentions=outputs.attentions,
                hidden_states=outputs.hidden_states,
            )

        return sequences


__all__ = ["GraniteSpeech5ForCTC", "GraniteSpeech5Encoder", "GraniteSpeech5PreTrainedModel"]