from __future__ import annotations import math from typing import Optional import numpy as np import torch import torch.nn.functional as F from torch import nn from transformers.activations import ACT2FN from transformers.modeling_outputs import BaseModelOutput from transformers.modeling_utils import PreTrainedModel from .qwen3_asr_audio_config import Qwen3ASRAudioEncoderConfig def _get_feat_extract_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor: input_lengths = torch.clamp(input_lengths.long(), min=1) input_lengths_leave = input_lengths % 100 feat_lengths = (input_lengths_leave - 1) // 2 + 1 return ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 def _block_diagonal_attention_mask(hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: seq_length = int(hidden_states.shape[0]) mask = torch.full( (1, 1, seq_length, seq_length), torch.finfo(hidden_states.dtype).min, dtype=hidden_states.dtype, device=hidden_states.device, ) for i in range(1, int(cu_seqlens.numel())): start = int(cu_seqlens[i - 1].item()) end = int(cu_seqlens[i].item()) mask[..., start:end, start:end] = 0 return mask class Qwen3ASRAudioAttention(nn.Module): def __init__(self, config: Qwen3ASRAudioEncoderConfig): super().__init__() self.embed_dim = int(config.d_model) self.num_heads = int(config.encoder_attention_heads) self.head_dim = self.embed_dim // self.num_heads if self.head_dim * self.num_heads != self.embed_dim: raise ValueError("d_model must be divisible by encoder_attention_heads") self.scaling = self.head_dim**-0.5 self.attention_dropout = float(config.attention_dropout) self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: del attention_mask seq_length = int(hidden_states.size(0)) query = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, self.head_dim).transpose(0, 1) key = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, self.head_dim).transpose(0, 1) value = self.v_proj(hidden_states).reshape(seq_length, self.num_heads, self.head_dim).transpose(0, 1) query = query.unsqueeze(0) key = key.unsqueeze(0) value = value.unsqueeze(0) mask = _block_diagonal_attention_mask(hidden_states, cu_seqlens) attn_weights = torch.matmul(query, key.transpose(2, 3)) * self.scaling attn_weights = attn_weights + mask attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(seq_length, self.embed_dim) return self.out_proj(attn_output) class Qwen3ASRAudioEncoderLayer(nn.Module): def __init__(self, config: Qwen3ASRAudioEncoderConfig): super().__init__() self.embed_dim = int(config.d_model) self.self_attn = Qwen3ASRAudioAttention(config) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.activation_fn = ACT2FN[config.activation_function] self.fc1 = nn.Linear(self.embed_dim, int(config.encoder_ffn_dim)) self.fc2 = nn.Linear(int(config.encoder_ffn_dim), self.embed_dim) self.final_layer_norm = nn.LayerNorm(self.embed_dim) def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> tuple[torch.Tensor]: residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states = self.self_attn(hidden_states=hidden_states, cu_seqlens=cu_seqlens) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16: clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) return (hidden_states,) class SinusoidsPositionEmbedding(nn.Module): def __init__(self, length: int, channels: int, max_timescale: int = 10000): super().__init__() if channels % 2 != 0: raise ValueError("SinusoidsPositionEmbedding requires an even channel count") log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float()) scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] self.register_buffer( "positional_embedding", torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1), persistent=False, ) def forward(self, seqlen: int): return self.positional_embedding[:seqlen, :] class Qwen3ASRAudioEncoder(PreTrainedModel): config_class = Qwen3ASRAudioEncoderConfig main_input_name = "input_features" _no_split_modules = ["Qwen3ASRAudioEncoderLayer"] def __init__(self, config: Qwen3ASRAudioEncoderConfig): super().__init__(config) embed_dim = int(config.d_model) self.dropout = float(config.dropout) self.num_mel_bins = int(config.num_mel_bins) self.max_source_positions = int(config.max_source_positions) self.embed_scale = math.sqrt(embed_dim) if bool(config.scale_embedding) else 1.0 self.n_window = int(config.n_window) self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim) self.layers = nn.ModuleList([Qwen3ASRAudioEncoderLayer(config) for _ in range(int(config.encoder_layers))]) self.ln_post = nn.LayerNorm(embed_dim) self.gradient_checkpointing = False self.conv2d1 = nn.Conv2d(1, int(config.downsample_hidden_size), 3, 2, padding=1) self.conv2d2 = nn.Conv2d(int(config.downsample_hidden_size), int(config.downsample_hidden_size), 3, 2, padding=1) self.conv2d3 = nn.Conv2d(int(config.downsample_hidden_size), int(config.downsample_hidden_size), 3, 2, padding=1) conv_freq = ((((int(config.num_mel_bins) + 1) // 2 + 1) // 2 + 1) // 2) self.conv_out = nn.Linear(int(config.downsample_hidden_size) * conv_freq, embed_dim, bias=False) self.proj1 = nn.Linear(embed_dim, embed_dim) self.act = ACT2FN[config.activation_function] self.proj2 = nn.Linear(embed_dim, int(config.output_dim)) self.n_window_infer = int(config.n_window_infer) self.conv_chunksize = int(config.conv_chunksize) self.post_init() def _freeze_parameters(self): for param in self.parameters(): param.requires_grad = False self._requires_grad = False def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: return _block_diagonal_attention_mask(inputs_tensor, cu_seqlens) def forward( self, input_features: torch.Tensor, feature_lens: Optional[torch.Tensor] = None, aftercnn_lens: Optional[torch.Tensor] = None, ): if feature_lens is None: feature_lens = torch.tensor([input_features.shape[-1]], dtype=torch.long, device=input_features.device) feature_lens = feature_lens.to(device=input_features.device, dtype=torch.long) if aftercnn_lens is None: aftercnn_lens = _get_feat_extract_output_lengths(feature_lens) aftercnn_lens = aftercnn_lens.to(device=input_features.device, dtype=torch.long) chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long() chunk_lengths = torch.tensor( [self.n_window * 2] * int(chunk_num.sum().item()), dtype=torch.long, device=feature_lens.device, ) tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:] chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2) chunk_lengths[chunk_lengths == 0] = self.n_window * 2 chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0) padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2) feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths) padded_mask_after_cnn = nn.utils.rnn.pad_sequence( [torch.ones(int(length.item()), dtype=torch.bool, device=padded_feature.device) for length in feature_lens_after_cnn], batch_first=True, ) padded_feature = padded_feature.unsqueeze(1) padded_embeds = [] for chunk in padded_feature.split(self.conv_chunksize, dim=0): padded_embed = F.gelu(self.conv2d1(chunk)) padded_embed = F.gelu(self.conv2d2(padded_embed)) padded_embed = F.gelu(self.conv2d3(padded_embed)) padded_embeds.append(padded_embed) padded_embed = torch.cat(padded_embeds, dim=0) bsz, channels, freq, time = padded_embed.size() padded_embed = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(bsz, time, channels * freq)) positional_embedding = ( self.positional_embedding.positional_embedding[: padded_embed.shape[1], :] .unsqueeze(0) .to(padded_embed.dtype) ) padded_embed = padded_embed + positional_embedding hidden_states = padded_embed[padded_mask_after_cnn] cu_chunk_lens = [0] window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2)) for cnn_len in aftercnn_lens: cnn_len_int = int(cnn_len.item()) cu_chunk_lens += [window_aftercnn] * (cnn_len_int // window_aftercnn) remainder = cnn_len_int % window_aftercnn if remainder != 0: cu_chunk_lens += [remainder] cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32) for encoder_layer in self.layers: hidden_states = encoder_layer(hidden_states, cu_seqlens)[0] hidden_states = self.ln_post(hidden_states) hidden_states = self.proj1(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.proj2(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) __all__ = ["Qwen3ASRAudioEncoder", "Qwen3ASRAudioEncoderLayer", "Qwen3ASRAudioAttention"]