# modeling_mini_whisper.py import math from collections.abc import Callable import torch import torch.nn as nn from torch.nn import functional as F from transformers import PreTrainedModel, GenerationMixin from transformers.modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, Seq2SeqLMOutput, Seq2SeqModelOutput, ) from transformers.activations import ACT2FN from transformers.utils import logging from configuration_mini_whisper import MiniWhisperConfig logger = logging.get_logger(__name__) _HIDDEN_STATES_START_POSITION = 1 def sinusoids(length: int, channels: int, max_timescale: float = 10000) -> torch.Tensor: """Returns sinusoids for positional embedding""" if channels % 2 != 0: raise ValueError(f"Number of channels has to be divisible by 2, got {channels} channels.") log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1) inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2)) scaled_time = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1) return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1) def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """Shift input ids one token to the right.""" shifted_input_ids = input_ids.new_zeros(input_ids.shape) shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() shifted_input_ids[:, 0] = decoder_start_token_id if pad_token_id is None: raise ValueError("pad_token_id has to be defined.") shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) return shifted_input_ids def _compute_mask_indices( shape: tuple[int, int], mask_prob: float, mask_length: int, attention_mask: torch.LongTensor | None = None, min_masks: int = 0, ): """Computes random mask spans for SpecAugment""" import numpy as np batch_size, sequence_length = shape if mask_length < 1: raise ValueError(f"mask_length has to be bigger than 0.") if mask_length > sequence_length: raise ValueError(f"mask_length has to be smaller than sequence_length.") epsilon = np.random.rand(1).item() def compute_num_masked_span(input_length): num_masked_span = int(mask_prob * input_length / mask_length + epsilon) num_masked_span = max(num_masked_span, min_masks) if num_masked_span * mask_length > sequence_length: num_masked_span = sequence_length // mask_length if input_length - (mask_length - 1) < num_masked_span: num_masked_span = max(input_length - (mask_length - 1), 0) return num_masked_span input_lengths = ( attention_mask.detach().sum(-1).tolist() if attention_mask is not None else [sequence_length for _ in range(batch_size)] ) spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool) spec_aug_mask_idxs = [] max_num_masked_span = compute_num_masked_span(sequence_length) if max_num_masked_span == 0: return spec_aug_mask for input_length in input_lengths: num_masked_span = compute_num_masked_span(input_length) spec_aug_mask_idx = np.random.choice( np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False ) if len(spec_aug_mask_idx) == 0: dummy_mask_idx = sequence_length - 1 else: dummy_mask_idx = spec_aug_mask_idx[0] spec_aug_mask_idx = np.concatenate( [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx] ) spec_aug_mask_idxs.append(spec_aug_mask_idx) spec_aug_mask_idxs = np.array(spec_aug_mask_idxs) spec_aug_mask_idxs = np.broadcast_to( spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length) ) spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length) offsets = np.arange(mask_length)[None, None, :] offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape( batch_size, max_num_masked_span * mask_length ) spec_aug_mask_idxs = spec_aug_mask_idxs + offsets if spec_aug_mask_idxs.max() > sequence_length - 1: spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1) return spec_aug_mask class WhisperPositionalEmbedding(nn.Module): def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None): super().__init__() self.num_positions = num_positions self.embedding_dim = embedding_dim self.weight = nn.Parameter(torch.zeros(num_positions, embedding_dim)) self._init_weights() def _init_weights(self): with torch.no_grad(): self.weight.copy_(sinusoids(self.num_positions, self.embedding_dim)) def forward(self, input_ids, past_key_values_length=0, position_ids=None): if position_ids is None: return self.weight[past_key_values_length : past_key_values_length + input_ids.shape[1]] else: return self.weight[position_ids] class MiniWhisperAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__( self, embed_dim: int, num_heads: int, dropout: float = 0.0, is_decoder: bool = False, bias: bool = True, is_causal: bool = False, layer_idx: int | None = None, config: MiniWhisperConfig | None = None, ): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.dropout = dropout self.head_dim = embed_dim // num_heads self.config = config if (self.head_dim * num_heads) != self.embed_dim: raise ValueError(f"embed_dim must be divisible by num_heads") self.scaling = self.head_dim ** -0.5 self.is_decoder = is_decoder self.is_causal = is_causal self.layer_idx = layer_idx self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias) self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) def forward( self, hidden_states: torch.Tensor, key_value_states: torch.Tensor | None = None, past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, output_attentions: bool = False, ): """Input shape: Batch x Time x Channel""" is_cross_attention = key_value_states is not None input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = (self.q_proj(hidden_states) * self.scaling).view(hidden_shape).transpose(1, 2).contiguous() if past_key_value is not None and not is_cross_attention: key_states = past_key_value[0] value_states = past_key_value[1] else: current_states = key_value_states if key_value_states is not None else hidden_states kv_shape = (input_shape[0], -1, self.num_heads, self.head_dim) key_states = self.k_proj(current_states).view(kv_shape).transpose(1, 2).contiguous() value_states = self.v_proj(current_states).view(kv_shape).transpose(1, 2).contiguous() if past_key_value is not None and is_cross_attention: key_states = torch.cat([past_key_value[0], key_states], dim=2) value_states = torch.cat([past_key_value[1], value_states], dim=2) attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * 1.0 if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = F.softmax(attn_weights, dim=-1) attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.out_proj(attn_output) return attn_output, None class MiniWhisperEncoderLayer(nn.Module): def __init__(self, config: MiniWhisperConfig): super().__init__() self.embed_dim = config.d_model self.self_attn = MiniWhisperAttention( embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads, dropout=config.attention_dropout, is_decoder=False, config=config, ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) self.final_layer_norm = nn.LayerNorm(self.embed_dim) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, ): residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, ) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = F.dropout(hidden_states, p=self.activation_dropout, training=self.training) hidden_states = self.fc2(hidden_states) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) 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 MiniWhisperDecoderLayer(nn.Module): def __init__(self, config: MiniWhisperConfig, layer_idx: int | None = None): super().__init__() self.embed_dim = config.d_model self.self_attn = MiniWhisperAttention( embed_dim=self.embed_dim, num_heads=config.decoder_attention_heads, dropout=config.attention_dropout, is_decoder=True, is_causal=True, layer_idx=layer_idx, config=config, ) self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.encoder_attn = MiniWhisperAttention( self.embed_dim, config.decoder_attention_heads, dropout=config.attention_dropout, is_decoder=True, layer_idx=layer_idx, config=config, ) self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim) self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim) self.final_layer_norm = nn.LayerNorm(self.embed_dim) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, past_key_value: tuple[tuple[torch.Tensor, torch.Tensor]] | None = None, use_cache: bool | None = True, ): residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) self_attn_past = past_key_value[0] if past_key_value is not None else None hidden_states, _ = self.self_attn( hidden_states, past_key_value=self_attn_past, attention_mask=attention_mask, ) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states if encoder_hidden_states is not None: residual = hidden_states hidden_states = self.encoder_attn_layer_norm(hidden_states) cross_attn_past = past_key_value[1] if past_key_value is not None else None hidden_states, _ = self.encoder_attn( hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, past_key_value=cross_attn_past, ) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = F.dropout(hidden_states, p=self.activation_dropout, training=self.training) hidden_states = self.fc2(hidden_states) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states return hidden_states, None class MiniWhisperEncoder(PreTrainedModel): config_class = MiniWhisperConfig main_input_name = "input_features" def __init__(self, config: MiniWhisperConfig): super().__init__(config) self.dropout = config.dropout self.layerdrop = config.encoder_layerdrop embed_dim = config.d_model self.num_mel_bins = config.num_mel_bins self.padding_idx = config.pad_token_id self.max_source_positions = config.max_source_positions self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1) self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1) self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim) self.embed_positions.requires_grad_(False) self.layers = nn.ModuleList([MiniWhisperEncoderLayer(config) for _ in range(config.encoder_layers)]) self.layer_norm = nn.LayerNorm(config.d_model) self.gradient_checkpointing = False self.post_init() def _init_weights(self, module): if isinstance(module, MiniWhisperEncoder): with torch.no_grad(): module.embed_positions.weight.copy_(sinusoids(*module.embed_positions.weight.shape)) def forward( self, input_features, attention_mask=None, **kwargs, ): x = F.gelu(self.conv1(input_features)) x = F.gelu(self.conv2(x)) x = x.permute(0, 2, 1) seq_len = x.shape[1] positions = torch.arange(seq_len, device=x.device) hidden_states = x + self.embed_positions(positions) hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) for idx, encoder_layer in enumerate(self.layers): if self.training: dropout_probability = torch.rand([]) if dropout_probability < self.layerdrop: continue hidden_states = encoder_layer(hidden_states, None) hidden_states = self.layer_norm(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): input_lengths = (input_lengths - 1) // 2 + 1 return input_lengths class MiniWhisperDecoder(PreTrainedModel): config_class = MiniWhisperConfig main_input_name = "input_ids" def __init__(self, config: MiniWhisperConfig): super().__init__(config) self.dropout = config.dropout self.layerdrop = config.decoder_layerdrop self.padding_idx = config.pad_token_id self.max_target_positions = config.max_target_positions self.max_source_positions = config.max_source_positions self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx) self.embed_positions = WhisperPositionalEmbedding(self.max_target_positions, config.d_model) self.layers = nn.ModuleList( [MiniWhisperDecoderLayer(config, layer_idx) for layer_idx in range(config.decoder_layers)] ) self.layer_norm = nn.LayerNorm(config.d_model) self.gradient_checkpointing = False self.post_init() def forward( self, input_ids=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_values=None, inputs_embeds=None, position_ids=None, use_cache=None, **kwargs, ): if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") elif input_ids is not None: inputs_embeds = self.embed_tokens(input_ids) elif inputs_embeds is None: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") past_key_values_length = 0 if past_key_values is not None and len(past_key_values) > 0: if past_key_values[0] is not None and past_key_values[0][0] is not None: past_key_values_length = past_key_values[0][0].shape[2] if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_key_values_length position_ids = position_ids.unsqueeze(0).repeat(inputs_embeds.shape[0], 1) if input_ids is not None: positions = self.embed_positions(input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids) else: positions = self.embed_positions(inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids) hidden_states = inputs_embeds + positions hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training) seq_len = inputs_embeds.shape[1] causal_mask = torch.triu( torch.ones(seq_len, seq_len, device=hidden_states.device) * float('-inf'), diagonal=1 ).unsqueeze(0).unsqueeze(0) if attention_mask is not None: padding_mask = (1 - attention_mask) * float('-inf') padding_mask = padding_mask.unsqueeze(1).unsqueeze(2) attention_mask = causal_mask + padding_mask else: attention_mask = causal_mask for idx, decoder_layer in enumerate(self.layers): if self.training: dropout_probability = torch.rand([]) if dropout_probability < self.layerdrop: continue layer_past = past_key_values[idx] if past_key_values is not None and idx < len(past_key_values) else None hidden_states, _ = decoder_layer( hidden_states, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, past_key_value=layer_past, use_cache=use_cache, ) hidden_states = self.layer_norm(hidden_states) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=None, ) class MiniWhisperModel(PreTrainedModel): config_class = MiniWhisperConfig def __init__(self, config: MiniWhisperConfig): super().__init__(config) self.encoder = MiniWhisperEncoder(config) self.decoder = MiniWhisperDecoder(config) self.post_init() def get_input_embeddings(self): return self.decoder.embed_tokens def set_input_embeddings(self, value): self.decoder.embed_tokens = value def _mask_input_features( self, input_features: torch.FloatTensor, attention_mask: torch.LongTensor | None = None, ): """Apply SpecAugment""" import numpy as np if not getattr(self.config, "apply_spec_augment", True): return input_features batch_size, hidden_size, sequence_length = input_features.size() if self.config.mask_time_prob > 0 and self.training: mask_time_indices = _compute_mask_indices( (batch_size, sequence_length), mask_prob=self.config.mask_time_prob, mask_length=self.config.mask_time_length, attention_mask=attention_mask, min_masks=self.config.mask_time_min_masks, ) mask_time_indices = torch.tensor(mask_time_indices, device=input_features.device, dtype=torch.bool) mask_time_indices = mask_time_indices[:, None].expand(-1, hidden_size, -1) input_features[mask_time_indices] = 0 if self.config.mask_feature_prob > 0 and self.training: mask_feature_indices = _compute_mask_indices( (batch_size, hidden_size), mask_prob=self.config.mask_feature_prob, mask_length=self.config.mask_feature_length, min_masks=self.config.mask_feature_min_masks, ) mask_feature_indices = torch.tensor(mask_feature_indices, device=input_features.device, dtype=torch.bool) input_features[mask_feature_indices] = 0 return input_features def forward( self, input_features: torch.FloatTensor | None = None, attention_mask: torch.LongTensor | None = None, decoder_input_ids: torch.LongTensor | None = None, decoder_attention_mask: torch.LongTensor | None = None, encoder_outputs: tuple[torch.FloatTensor] | None = None, past_key_values: tuple[tuple[torch.FloatTensor]] | None = None, decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, decoder_position_ids: tuple[torch.LongTensor] | None = None, use_cache: bool | None = None, **kwargs, ): if encoder_outputs is None: input_features = self._mask_input_features(input_features, attention_mask=attention_mask) encoder_outputs = self.encoder(input_features, **kwargs) elif not isinstance(encoder_outputs, BaseModelOutput): encoder_outputs = BaseModelOutput( last_hidden_state=encoder_outputs[0], hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, ) decoder_outputs = self.decoder( input_ids=decoder_input_ids, attention_mask=decoder_attention_mask, encoder_hidden_states=encoder_outputs.last_hidden_state, past_key_values=past_key_values, inputs_embeds=decoder_inputs_embeds, position_ids=decoder_position_ids, use_cache=use_cache, **kwargs, ) return Seq2SeqModelOutput( last_hidden_state=decoder_outputs.last_hidden_state, past_key_values=decoder_outputs.past_key_values, decoder_hidden_states=decoder_outputs.hidden_states, decoder_attentions=decoder_outputs.attentions, cross_attentions=decoder_outputs.cross_attentions, encoder_last_hidden_state=encoder_outputs.last_hidden_state, encoder_hidden_states=encoder_outputs.hidden_states, encoder_attentions=encoder_outputs.attentions, ) class MiniWhisperForConditionalGeneration(PreTrainedModel, GenerationMixin): config_class = MiniWhisperConfig base_model_prefix = "model" _tied_weights_keys = ["proj_out.weight"] def __init__(self, config: MiniWhisperConfig): super().__init__(config) self.model = MiniWhisperModel(config) self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False) self.max_target_positions = config.max_target_positions self.post_init() def get_output_embeddings(self): return self.proj_out def set_output_embeddings(self, new_embeddings): self.proj_out = new_embeddings def get_input_embeddings(self): return self.model.get_input_embeddings() def forward( self, input_features: torch.FloatTensor | None = None, attention_mask: torch.LongTensor | None = None, decoder_input_ids: torch.LongTensor | None = None, decoder_attention_mask: torch.LongTensor | None = None, encoder_outputs: tuple[torch.FloatTensor] | None = None, past_key_values: tuple[tuple[torch.FloatTensor]] | None = None, decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None, decoder_position_ids: tuple[torch.LongTensor] | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, **kwargs, ): if labels is not None: if labels.shape[1] > self.max_target_positions: raise ValueError(f"Labels' sequence length {labels.shape[1]} cannot exceed max_target_positions") if decoder_input_ids is None and decoder_inputs_embeds is None: decoder_input_ids = shift_tokens_right( labels, self.config.pad_token_id, self.config.decoder_start_token_id ) outputs = self.model( input_features, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_outputs, decoder_attention_mask=decoder_attention_mask, past_key_values=past_key_values, decoder_inputs_embeds=decoder_inputs_embeds, decoder_position_ids=decoder_position_ids, use_cache=use_cache, **kwargs, ) lm_logits = self.proj_out(outputs.last_hidden_state) loss = None if labels is not None: loss_fct = nn.CrossEntropyLoss() labels = labels.to(lm_logits.device) loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1)) return Seq2SeqLMOutput( loss=loss, logits=lm_logits, past_key_values=outputs.past_key_values, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions, ) def generate( self, input_features, generation_config=None, max_new_tokens=50, **kwargs, ): self.eval() batch_size = input_features.shape[0] device = input_features.device decoder_input_ids = torch.full( (batch_size, 1), self.config.decoder_start_token_id, dtype=torch.long, device=device ) encoder_outputs = self.model.encoder(input_features) for _ in range(max_new_tokens): outputs = self.forward( decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_outputs, use_cache=True, ) next_token_logits = outputs.logits[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) decoder_input_ids = torch.cat([decoder_input_ids, next_token], dim=1) if (next_token == self.config.eos_token_id).all(): break return decoder_input_ids def prepare_inputs_for_generation( self, decoder_input_ids, past_key_values=None, encoder_outputs=None, attention_mask=None, use_cache=None, **kwargs, ): return { "decoder_input_ids": decoder_input_ids, "past_key_values": past_key_values, "encoder_outputs": encoder_outputs, "attention_mask": attention_mask, "use_cache": use_cache, } def get_encoder(self): return self.model.get_encoder() __all__ = [ "MiniWhisperConfig", "MiniWhisperModel", "MiniWhisperForConditionalGeneration", "MiniWhisperEncoder", "MiniWhisperDecoder", "MiniWhisperEncoderLayer", "MiniWhisperDecoderLayer", "MiniWhisperAttention", "WhisperPositionalEmbedding", ]