Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/moonshine/modular_moonshine.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_moonshine.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# Copyright 2025 The HuggingFace Inc. team. All rights reserved.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20 21from typing import Callable, Optional, Union22 23import numpy as np24import torch25import torch.nn as nn26 27from transformers.utils.generic import OutputRecorder, check_model_inputs28 29from ...activations import ACT2FN30from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache31from ...generation import GenerationMixin32from ...masking_utils import create_causal_mask33from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa34from ...modeling_flash_attention_utils import FlashAttentionKwargs35from ...modeling_layers import GradientCheckpointingLayer36from ...modeling_outputs import (37 BaseModelOutput,38 BaseModelOutputWithPast,39 BaseModelOutputWithPastAndCrossAttentions,40 Seq2SeqLMOutput,41 Seq2SeqModelOutput,42)43from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update44from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel45from ...processing_utils import Unpack46from ...utils import TransformersKwargs, auto_docstring, can_return_tuple47from ...utils.deprecation import deprecate_kwarg48from .configuration_moonshine import MoonshineConfig49 50 51class MoonshineEncoderMLP(nn.Module):52 def __init__(self, config, hidden_act):53 super().__init__()54 self.config = config55 self.activation_fn = ACT2FN[hidden_act]56 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)57 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)58 59 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:60 hidden_states = self.fc1(hidden_states)61 hidden_states = self.activation_fn(hidden_states)62 hidden_states = self.fc2(hidden_states)63 return hidden_states64 65 66class MoonshineDecoderMLP(nn.Module):67 def __init__(self, config, hidden_act):68 super().__init__()69 self.config = config70 self.activation_fn = ACT2FN[hidden_act]71 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size * 2)72 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)73 74 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:75 hidden_states = self.fc1(hidden_states)76 hidden_states, gate = hidden_states.chunk(2, dim=-1)77 hidden_states = self.activation_fn(gate) * hidden_states78 hidden_states = self.fc2(hidden_states)79 return hidden_states80 81 82def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:83 """84 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,85 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)86 """87 batch, num_key_value_heads, slen, head_dim = hidden_states.shape88 if n_rep == 1:89 return hidden_states90 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)91 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)92 93 94def eager_attention_forward(95 module: nn.Module,96 query: torch.Tensor,97 key: torch.Tensor,98 value: torch.Tensor,99 attention_mask: Optional[torch.Tensor],100 scaling: float,101 dropout: float = 0.0,102 **kwargs: Unpack[TransformersKwargs],103):104 key_states = repeat_kv(key, module.num_key_value_groups)105 value_states = repeat_kv(value, module.num_key_value_groups)106 107 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling108 if attention_mask is not None:109 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]110 attn_weights = attn_weights + causal_mask111 112 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)113 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)114 attn_output = torch.matmul(attn_weights, value_states)115 attn_output = attn_output.transpose(1, 2).contiguous()116 117 return attn_output, attn_weights118 119 120def rotate_half(x):121 """Rotates half the hidden dims of the input."""122 x1 = x[..., 0::2]123 x2 = x[..., 1::2]124 return torch.stack((-x2, x1), dim=-1).flatten(-2)125 126 127def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):128 """Applies Rotary Position Embedding to the query and key tensors.129 130 Args:131 q (`torch.Tensor`): The query tensor.132 k (`torch.Tensor`): The key tensor.133 cos (`torch.Tensor`): The cosine part of the rotary embedding.134 sin (`torch.Tensor`): The sine part of the rotary embedding.135 position_ids (`torch.Tensor`, *optional*):136 Deprecated and unused.137 unsqueeze_dim (`int`, *optional*, defaults to 1):138 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and139 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note140 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and141 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes142 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have143 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.144 Returns:145 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.146 """147 cos = cos.unsqueeze(unsqueeze_dim)148 sin = sin.unsqueeze(unsqueeze_dim)149 150 # Interleave them instead of usual shape151 cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)152 sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)153 154 # Keep half or full tensor for later concatenation155 rotary_dim = cos.shape[-1]156 q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]157 k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]158 159 # Apply rotary embeddings on the first half or full tensor160 q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)161 k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)162 163 # Concatenate back to full shape164 q_embed = torch.cat([q_embed, q_pass], dim=-1)165 k_embed = torch.cat([k_embed, k_pass], dim=-1)166 return q_embed, k_embed167 168 169class MoonshineAttention(nn.Module):170 """Multi-headed attention from 'Attention Is All You Need' paper"""171 172 def __init__(173 self,174 config: MoonshineConfig,175 layer_idx: int,176 is_causal: bool,177 num_attention_heads: int,178 num_key_value_heads: int,179 ):180 super().__init__()181 config.update({"num_attention_heads": num_attention_heads, "num_key_value_heads": num_key_value_heads})182 self.config = config183 self.layer_idx = layer_idx184 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)185 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads186 self.scaling = self.head_dim**-0.5187 self.attention_dropout = config.attention_dropout188 self.is_causal = is_causal189 190 self.q_proj = nn.Linear(191 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias192 )193 self.k_proj = nn.Linear(194 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias195 )196 self.v_proj = nn.Linear(197 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias198 )199 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)200 201 # Pad head dimension to the next specified multiple.202 if self.config.pad_head_dim_to_multiple_of is not None:203 target_multiple = self.config.pad_head_dim_to_multiple_of204 target_head_dim = target_multiple * ((self.head_dim + target_multiple - 1) // target_multiple)205 self.head_dim_padding = target_head_dim - self.head_dim206 else:207 self.head_dim_padding = 0208 209 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")210 def forward(211 self,212 hidden_states: torch.Tensor,213 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,214 attention_mask: Optional[torch.Tensor] = None,215 past_key_values: Optional[Cache] = None,216 cache_position: Optional[torch.LongTensor] = None,217 key_value_states: Optional[torch.Tensor] = None,218 **kwargs: Unpack[FlashAttentionKwargs],219 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:220 bsz, q_len = hidden_states.shape[:-1]221 222 query_states = (223 self.q_proj(hidden_states).view(bsz, q_len, self.config.num_key_value_heads, self.head_dim).transpose(1, 2)224 )225 226 is_cross_attention = key_value_states is not None227 if past_key_values is not None:228 is_updated = past_key_values.is_updated.get(self.layer_idx)229 if is_cross_attention:230 # after the first generated id, we can subsequently re-use all key/value_states from cache231 past_key_values.is_updated[self.layer_idx] = True232 past_key_values = past_key_values.cross_attention_cache233 else:234 past_key_values = past_key_values.self_attention_cache235 236 # use key_value_states if cross attention237 current_states = key_value_states if key_value_states is not None else hidden_states238 if is_cross_attention and past_key_values and is_updated:239 key_states = past_key_values.layers[self.layer_idx].keys240 value_states = past_key_values.layers[self.layer_idx].values241 else:242 key_states = (243 self.k_proj(current_states)244 .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)245 .transpose(1, 2)246 )247 value_states = (248 self.v_proj(current_states)249 .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)250 .transpose(1, 2)251 )252 if is_cross_attention and past_key_values is not None:253 key_states, value_states = past_key_values.update(254 key_states, value_states, self.layer_idx, {"cache_position": cache_position}255 )256 257 if not is_cross_attention:258 cos, sin = position_embeddings259 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)260 261 if past_key_values is not None:262 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}263 key_states, value_states = past_key_values.update(264 key_states, value_states, self.layer_idx, cache_kwargs265 )266 267 attention_interface: Callable = eager_attention_forward268 if self.config._attn_implementation != "eager":269 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]270 271 is_causal = self.is_causal and attention_mask is None and q_len > 1272 273 if self.head_dim_padding > 0:274 query_states = torch.nn.functional.pad(query_states, (0, self.head_dim_padding))275 key_states = torch.nn.functional.pad(key_states, (0, self.head_dim_padding))276 value_states = torch.nn.functional.pad(value_states, (0, self.head_dim_padding))277 278 attn_output, attn_weights = attention_interface(279 self,280 query_states,281 key_states,282 value_states,283 attention_mask,284 dropout=0.0 if not self.training else self.attention_dropout,285 scaling=self.scaling,286 is_causal=is_causal,287 **kwargs,288 )289 290 if self.head_dim_padding > 0:291 attn_output = attn_output[..., : -self.head_dim_padding]292 293 attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()294 attn_output = self.o_proj(attn_output)295 return attn_output, attn_weights296 297 298class MoonshineRotaryEmbedding(nn.Module):299 inv_freq: torch.Tensor # fix linting for `register_buffer`300 301 def __init__(self, config: MoonshineConfig, device=None):302 super().__init__()303 # BC: "rope_type" was originally "type"304 if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):305 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))306 else:307 self.rope_type = "default"308 self.max_seq_len_cached = config.max_position_embeddings309 self.original_max_seq_len = config.max_position_embeddings310 311 self.config = config312 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]313 314 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)315 self.register_buffer("inv_freq", inv_freq, persistent=False)316 self.original_inv_freq = self.inv_freq317 318 @torch.no_grad()319 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)320 def forward(self, x, position_ids):321 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)322 position_ids_expanded = position_ids[:, None, :].float()323 324 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"325 with torch.autocast(device_type=device_type, enabled=False): # Force float32326 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)327 emb = torch.cat((freqs, freqs), dim=-1)328 cos = emb.cos() * self.attention_scaling329 sin = emb.sin() * self.attention_scaling330 331 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)332 333 334class MoonshineEncoderLayer(GradientCheckpointingLayer):335 def __init__(self, config: MoonshineConfig, layer_idx: int):336 super().__init__()337 self.hidden_size = config.hidden_size338 339 self.self_attn = MoonshineAttention(340 config=config,341 layer_idx=layer_idx,342 is_causal=False,343 num_attention_heads=config.encoder_num_attention_heads,344 num_key_value_heads=config.encoder_num_key_value_heads,345 )346 347 self.mlp = MoonshineEncoderMLP(config, config.encoder_hidden_act)348 self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)349 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)350 351 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")352 def forward(353 self,354 hidden_states: torch.Tensor,355 attention_mask: Optional[torch.Tensor] = None,356 position_ids: Optional[torch.LongTensor] = None,357 past_key_values: Optional[Cache] = None,358 use_cache: Optional[bool] = False,359 cache_position: Optional[torch.LongTensor] = None,360 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC361 **kwargs: Unpack[TransformersKwargs],362 ) -> torch.Tensor:363 residual = hidden_states364 hidden_states = self.input_layernorm(hidden_states)365 # Self Attention366 hidden_states, _ = self.self_attn(367 hidden_states=hidden_states,368 attention_mask=attention_mask,369 position_ids=position_ids,370 past_key_values=past_key_values,371 use_cache=use_cache,372 cache_position=cache_position,373 position_embeddings=position_embeddings,374 **kwargs,375 )376 hidden_states = residual + hidden_states377 378 # Fully Connected379 residual = hidden_states380 hidden_states = self.post_attention_layernorm(hidden_states)381 hidden_states = self.mlp(hidden_states)382 hidden_states = residual + hidden_states383 return hidden_states384 385 386class MoonshineDecoderLayer(GradientCheckpointingLayer):387 def __init__(self, config: MoonshineConfig, layer_idx: Optional[int] = None):388 super().__init__()389 self.hidden_size = config.hidden_size390 391 self.self_attn = MoonshineAttention(392 config=config,393 layer_idx=layer_idx,394 is_causal=True,395 num_attention_heads=config.decoder_num_attention_heads,396 num_key_value_heads=config.decoder_num_key_value_heads,397 )398 self.encoder_attn = MoonshineAttention(399 config=config,400 layer_idx=layer_idx,401 is_causal=False,402 num_attention_heads=config.decoder_num_attention_heads,403 num_key_value_heads=config.decoder_num_key_value_heads,404 )405 406 self.mlp = MoonshineDecoderMLP(config, config.decoder_hidden_act)407 self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)408 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)409 self.final_layernorm = nn.LayerNorm(config.hidden_size, bias=False)410 411 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")412 def forward(413 self,414 hidden_states: torch.Tensor,415 attention_mask: Optional[torch.Tensor] = None,416 encoder_hidden_states: Optional[torch.Tensor] = None,417 encoder_attention_mask: Optional[torch.Tensor] = None,418 position_ids: Optional[torch.LongTensor] = None,419 encoder_position_ids: Optional[torch.LongTensor] = None,420 past_key_values: Optional[Cache] = None,421 use_cache: Optional[bool] = False,422 cache_position: Optional[torch.LongTensor] = None,423 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,424 encoder_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,425 **kwargs: Unpack[TransformersKwargs],426 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:427 residual = hidden_states428 hidden_states = self.input_layernorm(hidden_states)429 430 hidden_states, _ = self.self_attn(431 hidden_states=hidden_states,432 attention_mask=attention_mask,433 position_ids=position_ids,434 past_key_values=past_key_values,435 use_cache=use_cache,436 cache_position=cache_position,437 position_embeddings=position_embeddings,438 **kwargs,439 )440 hidden_states = residual + hidden_states441 442 if encoder_hidden_states is not None:443 residual = hidden_states444 hidden_states = self.post_attention_layernorm(hidden_states)445 hidden_states, _ = self.encoder_attn(446 hidden_states=hidden_states,447 key_value_states=encoder_hidden_states,448 attention_mask=encoder_attention_mask,449 past_key_values=past_key_values,450 use_cache=use_cache,451 )452 hidden_states = residual + hidden_states453 454 residual = hidden_states455 hidden_states = self.final_layernorm(hidden_states)456 hidden_states = self.mlp(hidden_states)457 hidden_states = residual + hidden_states458 return hidden_states459 460 461@auto_docstring462class MoonshinePreTrainedModel(PreTrainedModel):463 config: MoonshineConfig464 base_model_prefix = "model"465 main_input_name = "input_values"466 supports_gradient_checkpointing = True467 _no_split_modules = ["MoonshineEncoderLayer", "MoonshineDecoderLayer"]468 _supports_flash_attn = True469 _supports_sdpa = True470 471 _can_compile_fullgraph = True472 # TODO arthur, how do we separate when it cross / self coming from different layer?473 474 def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):475 """476 Computes the output length of the convolutional layers477 """478 output_conv1_length = int((input_lengths - 127) / 64 + 1)479 output_conv2_length = int((output_conv1_length - 7) / 3 + 1)480 output_conv3_length = int((output_conv2_length - 3) / 2 + 1)481 482 return output_conv3_length483 484 485class MoonshineEncoder(MoonshinePreTrainedModel):486 """487 Transformer encoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MoonshineEncoderLayer`]488 489 Args:490 config: MoonshineConfig491 """492 493 main_input_name = "input_values"494 _can_record_outputs = {495 "attentions": MoonshineAttention,496 "hidden_states": MoonshineEncoderLayer,497 }498 499 def __init__(self, config: MoonshineConfig):500 super().__init__(config)501 self.config = config502 embed_dim = config.hidden_size503 504 self.conv1 = nn.Conv1d(1, embed_dim, kernel_size=127, stride=64, bias=False)505 self.conv2 = nn.Conv1d(embed_dim, 2 * embed_dim, kernel_size=7, stride=3)506 self.conv3 = nn.Conv1d(2 * embed_dim, embed_dim, kernel_size=3, stride=2)507 self.groupnorm = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=1e-5)508 self.rotary_emb = MoonshineRotaryEmbedding(config=config)509 510 self.layers = nn.ModuleList(511 [MoonshineEncoderLayer(config, idx) for idx in range(config.encoder_num_hidden_layers)]512 )513 self.layer_norm = nn.LayerNorm(embed_dim, bias=False)514 self.gradient_checkpointing = False515 self.post_init()516 517 def get_input_embeddings(self) -> nn.Module:518 return self.conv1519 520 def set_input_embeddings(self, value: nn.Module):521 self.conv1 = value522 523 @check_model_inputs()524 def forward(525 self,526 input_values: torch.FloatTensor,527 attention_mask: Optional[torch.Tensor] = None,528 **kwargs: Unpack[TransformersKwargs],529 ) -> BaseModelOutputWithPast:530 r"""531 Args:532 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):533 Float values of the raw speech waveform. Raw speech waveform can be534 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a535 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or536 the soundfile library (`pip install soundfile`). To prepare the array into537 `input_values`, the [`AutoFeatureExtractor`] should be used for padding538 and conversion into a tensor of type `torch.FloatTensor`.539 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):540 Mask to avoid performing attention on padding indices in `input_values`. Mask values selected in `[0, 1]`:541 - 1 for tokens that are **not masked**,542 - 0 for tokens that are **masked**.543 [What are attention masks?](../glossary#attention-mask)544 """545 input_values = input_values.unsqueeze(1)546 hidden_states = nn.functional.tanh(self.conv1(input_values))547 hidden_states = self.groupnorm(hidden_states)548 hidden_states = nn.functional.gelu(self.conv2(hidden_states))549 hidden_states = nn.functional.gelu(self.conv3(hidden_states))550 hidden_states = hidden_states.permute(0, 2, 1)551 552 # attention mask downsampling553 if attention_mask is not None:554 mask_len = self._get_feat_extract_output_lengths(attention_mask.shape[-1])555 downsample_stride = 64 * 3 * 2 # conv strides556 attention_mask = attention_mask[..., ::downsample_stride][..., :mask_len]557 if self.config._attn_implementation == "flash_attention_2":558 attention_mask = attention_mask if (attention_mask == 0.0).any() else None559 elif self.config._attn_implementation == "sdpa":560 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, hidden_states.dtype)561 else:562 attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)563 564 position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)565 position_embeddings = self.rotary_emb(hidden_states, position_ids)566 567 for encoder_layer in self.layers:568 hidden_states = encoder_layer(569 hidden_states,570 attention_mask=attention_mask,571 position_ids=position_ids,572 position_embeddings=position_embeddings,573 **kwargs,574 )575 576 hidden_states = self.layer_norm(hidden_states)577 578 return BaseModelOutputWithPast(579 last_hidden_state=hidden_states,580 )581 582 583@auto_docstring584class MoonshineDecoder(MoonshinePreTrainedModel):585 main_input_name = "input_ids"586 _can_record_outputs = {587 "attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="self_attn"),588 "hidden_states": MoonshineDecoderLayer,589 "cross_attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="encoder_attn"),590 }591 592 def __init__(self, config: MoonshineConfig):593 super().__init__(config)594 self.padding_idx = config.pad_token_id595 self.vocab_size = config.vocab_size596 597 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)598 self.layers = nn.ModuleList(599 [MoonshineDecoderLayer(config, idx) for idx in range(config.decoder_num_hidden_layers)]600 )601 self.norm = nn.LayerNorm(config.hidden_size, bias=False)602 self.rotary_emb = MoonshineRotaryEmbedding(config=config)603 self.gradient_checkpointing = False604 605 # Initialize weights and apply final processing606 self.post_init()607 608 @check_model_inputs()609 def forward(610 self,611 input_ids: Optional[torch.LongTensor] = None,612 attention_mask: Optional[torch.Tensor] = None,613 position_ids: Optional[torch.LongTensor] = None,614 past_key_values: Optional[Cache] = None,615 inputs_embeds: Optional[torch.FloatTensor] = None,616 use_cache: Optional[bool] = None,617 cache_position: Optional[torch.LongTensor] = None,618 encoder_hidden_states: Optional[torch.FloatTensor] = None,619 encoder_attention_mask: Optional[torch.Tensor] = None,620 **kwargs: Unpack[TransformersKwargs],621 ) -> Union[tuple, BaseModelOutputWithPast]:622 r"""623 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):624 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention625 of the decoder.626 encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):627 Mask to avoid performing attention on padding indices in `encoder_hidden_states`. Mask values selected in `[0, 1]`:628 - 1 for tokens that are **not masked**,629 - 0 for tokens that are **masked**.630 [What are attention masks?](../glossary#attention-mask)631 """632 if (input_ids is None) ^ (inputs_embeds is not None):633 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")634 635 if inputs_embeds is None:636 inputs_embeds = self.embed_tokens(input_ids)637 638 if use_cache and past_key_values is None:639 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))640 641 if cache_position is None:642 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0643 cache_position = torch.arange(644 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device645 )646 647 if position_ids is None:648 position_ids = cache_position.unsqueeze(0)649 650 causal_mask = create_causal_mask(651 config=self.config,652 input_embeds=inputs_embeds,653 attention_mask=attention_mask,654 cache_position=cache_position,655 past_key_values=past_key_values,656 position_ids=position_ids,657 )658 659 hidden_states = inputs_embeds660 position_embeddings = self.rotary_emb(hidden_states, position_ids)661 662 if encoder_attention_mask is not None:663 mask_len = encoder_hidden_states.shape[-2]664 downsample_stride = 64 * 3 * 2 # conv strides665 encoder_attention_mask = encoder_attention_mask[..., ::downsample_stride][..., :mask_len]666 if self.config._attn_implementation == "flash_attention_2":667 encoder_attention_mask = encoder_attention_mask if (encoder_attention_mask == 0.0).any() else None668 elif self.config._attn_implementation == "sdpa":669 encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(670 encoder_attention_mask, hidden_states.dtype, hidden_states.shape[-2]671 )672 else:673 encoder_attention_mask = _prepare_4d_attention_mask(674 encoder_attention_mask, hidden_states.dtype, hidden_states.shape[-2]675 )676 677 for decoder_layer in self.layers:678 hidden_states = decoder_layer(679 hidden_states,680 causal_mask,681 encoder_hidden_states, # as a positional argument for gradient checkpointing682 encoder_attention_mask=encoder_attention_mask,683 position_ids=position_ids,684 past_key_values=past_key_values,685 use_cache=use_cache,686 cache_position=cache_position,687 position_embeddings=position_embeddings,688 **kwargs,689 )690 691 hidden_states = self.norm(hidden_states)692 693 return BaseModelOutputWithPastAndCrossAttentions(694 last_hidden_state=hidden_states,695 past_key_values=past_key_values if use_cache else None,696 )697 698 699def _compute_mask_indices(700 shape: tuple[int, int],701 mask_prob: float,702 mask_length: int,703 attention_mask: Optional[torch.LongTensor] = None,704 min_masks: int = 0,705) -> np.ndarray:706 """707 Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for708 ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on709 CPU as part of the preprocessing during training.710 711 Args:712 shape: The shape for which to compute masks. This should be of a tuple of size 2 where713 the first element is the batch size and the second element is the length of the axis to span.714 mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of715 independently generated mask spans of length `mask_length` is computed by716 `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the717 actual percentage will be smaller.718 mask_length: size of the mask719 min_masks: minimum number of masked spans720 attention_mask: A (right-padded) attention mask which independently shortens the feature axis of721 each batch dimension.722 """723 batch_size, sequence_length = shape724 725 if mask_length < 1:726 raise ValueError("`mask_length` has to be bigger than 0.")727 728 if mask_length > sequence_length:729 raise ValueError(730 f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"731 f" and `sequence_length`: {sequence_length}`"732 )733 734 # epsilon is used for probabilistic rounding735 epsilon = np.random.rand(1).item()736 737 def compute_num_masked_span(input_length):738 """Given input length, compute how many spans should be masked"""739 num_masked_span = int(mask_prob * input_length / mask_length + epsilon)740 num_masked_span = max(num_masked_span, min_masks)741 742 # make sure num masked span <= sequence_length743 if num_masked_span * mask_length > sequence_length:744 num_masked_span = sequence_length // mask_length745 746 # make sure num_masked span is also <= input_length - (mask_length - 1)747 if input_length - (mask_length - 1) < num_masked_span:748 num_masked_span = max(input_length - (mask_length - 1), 0)749 750 return num_masked_span751 752 # compute number of masked spans in batch753 input_lengths = (754 attention_mask.detach().sum(-1).tolist()755 if attention_mask is not None756 else [sequence_length for _ in range(batch_size)]757 )758 759 # SpecAugment mask to fill760 spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)761 spec_aug_mask_idxs = []762 763 max_num_masked_span = compute_num_masked_span(sequence_length)764 765 if max_num_masked_span == 0:766 return spec_aug_mask767 768 for input_length in input_lengths:769 # compute num of masked spans for this input770 num_masked_span = compute_num_masked_span(input_length)771 772 # get random indices to mask773 spec_aug_mask_idx = np.random.choice(774 np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False775 )776 777 # pick first sampled index that will serve as a dummy index to pad vector778 # to ensure same dimension for all batches due to probabilistic rounding779 # Picking first sample just pads those vectors twice.780 if len(spec_aug_mask_idx) == 0:781 # this case can only happen if `input_length` is strictly smaller then782 # `sequence_length` in which case the last token has to be a padding783 # token which we can use as a dummy mask id784 dummy_mask_idx = sequence_length - 1785 else:786 dummy_mask_idx = spec_aug_mask_idx[0]787 788 spec_aug_mask_idx = np.concatenate(789 [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]790 )791 spec_aug_mask_idxs.append(spec_aug_mask_idx)792 793 spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)794 795 # expand masked indices to masked spans796 spec_aug_mask_idxs = np.broadcast_to(797 spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)798 )799 spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)800 801 # add offset to the starting indexes so that indexes now create a span802 offsets = np.arange(mask_length)[None, None, :]803 offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(804 batch_size, max_num_masked_span * mask_length805 )806 spec_aug_mask_idxs = spec_aug_mask_idxs + offsets807 808 # ensure that we cannot have indices larger than sequence_length809 if spec_aug_mask_idxs.max() > sequence_length - 1:810 spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1811 812 # scatter indices to mask813 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)814 815 return spec_aug_mask816 817 818@auto_docstring819class MoonshineModel(MoonshinePreTrainedModel):820 def __init__(self, config: MoonshineConfig):821 super().__init__(config)822 823 self.encoder = MoonshineEncoder(config)824 self.decoder = MoonshineDecoder(config)825 # Initialize weights and apply final processing826 self.post_init()827 828 def get_input_embeddings(self):829 return self.decoder.embed_tokens830 831 def set_input_embeddings(self, value):832 self.decoder.embed_tokens = value833 834 def get_encoder(self):835 return self.encoder836 837 def freeze_encoder(self):838 """839 Calling this function will disable the gradient computation for the Moonshine encoder so that its parameters will840 not be updated during training.841 """842 self.encoder._freeze_parameters()843 844 def _mask_input_features(845 self,846 input_features: torch.FloatTensor,847 attention_mask: Optional[torch.LongTensor] = None,848 ):849 """850 Masks extracted features along time axis and/or along feature axis according to851 [SpecAugment](https://huggingface.co/papers/1904.08779).852 """853 854 # `config.apply_spec_augment` can set masking to False855 if not getattr(self.config, "apply_spec_augment", True):856 return input_features857 858 # generate indices & apply SpecAugment along time axis859 batch_size, hidden_size, sequence_length = input_features.size()860 861 if self.config.mask_time_prob > 0 and self.training:862 # generate indices & apply SpecAugment along time axis863 mask_time_indices = _compute_mask_indices(864 (batch_size, sequence_length),865 mask_prob=self.config.mask_time_prob,866 mask_length=self.config.mask_time_length,867 attention_mask=attention_mask,868 min_masks=self.config.mask_time_min_masks,869 )870 mask_time_indices = torch.tensor(mask_time_indices, device=input_features.device, dtype=torch.bool)871 mask_time_indices = mask_time_indices[:, None].expand(-1, hidden_size, -1)872 input_features[mask_time_indices] = 0873 874 if self.config.mask_feature_prob > 0 and self.training:875 # generate indices & apply SpecAugment along feature axis876 mask_feature_indices = _compute_mask_indices(877 (batch_size, hidden_size),878 mask_prob=self.config.mask_feature_prob,879 mask_length=self.config.mask_feature_length,880 min_masks=self.config.mask_feature_min_masks,881 )882 mask_feature_indices = torch.tensor(mask_feature_indices, device=input_features.device, dtype=torch.bool)883 input_features[mask_feature_indices] = 0884 885 return input_features886 887 @can_return_tuple888 @auto_docstring889 def forward(890 self,891 input_values: Optional[torch.FloatTensor] = None,892 attention_mask: Optional[torch.LongTensor] = None,893 decoder_input_ids: Optional[torch.LongTensor] = None,894 decoder_attention_mask: Optional[torch.LongTensor] = None,895 encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,896 past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None,897 decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None,898 decoder_position_ids: Optional[tuple[torch.LongTensor]] = None,899 use_cache: Optional[bool] = None,900 cache_position: Optional[torch.LongTensor] = None,901 **kwargs: Unpack[TransformersKwargs],902 ) -> Seq2SeqModelOutput:903 r"""904 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):905 Float values of the raw speech waveform. Raw speech waveform can be906 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a907 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or908 the soundfile library (`pip install soundfile`). To prepare the array into909 `input_values`, the [`AutoFeatureExtractor`] should be used for padding910 and conversion into a tensor of type `torch.FloatTensor`.911 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):912 Indices of positions of each input sequence tokens in the position embeddings.913 Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`914 915 Example:916 917 ```python918 >>> import torch919 >>> from transformers import AutoFeatureExtractor, MoonshineModel920 >>> from datasets import load_dataset921 922 >>> model = MoonshineModel.from_pretrained("UsefulSensors/moonshine-tiny")923 >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/moonshine-tiny")924 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")925 >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")926 >>> input_values = inputs.input_values927 >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id928 >>> last_hidden_state = model(input_values, decoder_input_ids=decoder_input_ids).last_hidden_state929 >>> list(last_hidden_state.shape)930 [1, 2, 288]931 ```932 """933 if encoder_outputs is None:934 encoder_outputs: BaseModelOutput = self.encoder(input_values, attention_mask=attention_mask, **kwargs)935 936 decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(937 input_ids=decoder_input_ids,938 attention_mask=decoder_attention_mask,939 encoder_attention_mask=attention_mask,940 encoder_hidden_states=encoder_outputs.last_hidden_state,941 past_key_values=past_key_values,942 inputs_embeds=decoder_inputs_embeds,943 position_ids=decoder_position_ids,944 use_cache=use_cache,945 cache_position=cache_position,946 **kwargs,947 )948 949 return Seq2SeqModelOutput(950 last_hidden_state=decoder_outputs.last_hidden_state,951 past_key_values=decoder_outputs.past_key_values,952 decoder_hidden_states=decoder_outputs.hidden_states,953 decoder_attentions=decoder_outputs.attentions,954 cross_attentions=decoder_outputs.cross_attentions,955 encoder_last_hidden_state=encoder_outputs.last_hidden_state,956 encoder_hidden_states=encoder_outputs.hidden_states,957 encoder_attentions=encoder_outputs.attentions,958 )959 960 961def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):962 """963 Shift input ids one token to the right.964 """965 shifted_input_ids = input_ids.new_zeros(input_ids.shape)966 shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()967 shifted_input_ids[:, 0] = decoder_start_token_id968 969 if pad_token_id is None:970 raise ValueError("self.model.config.pad_token_id has to be defined.")971 # replace possible -100 values in labels by `pad_token_id`972 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)973 974 return shifted_input_ids975 976 977@auto_docstring(978 custom_intro="""979 The Moonshine Model with a language modeling head. Can be used for automatic speech recognition.980 """981)982class MoonshineForConditionalGeneration(MoonshinePreTrainedModel, GenerationMixin):983 _tied_weights_keys = ["proj_out.weight"]984 985 def __init__(self, config: MoonshineConfig):986 super().__init__(config)987 self.model = MoonshineModel(config)988 self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)989 990 # Initialize weights and apply final processing991 self.post_init()992 993 def get_encoder(self):994 return self.model.get_encoder()995 996 def get_decoder(self):997 return self.model.get_decoder()998 999 def get_output_embeddings(self):1000 return self.proj_out1001 1002 def set_output_embeddings(self, new_embeddings):1003 self.proj_out = new_embeddings1004 1005 def get_input_embeddings(self) -> nn.Module:1006 return self.model.get_input_embeddings()1007 1008 @can_return_tuple1009 @auto_docstring1010 def forward(1011 self,1012 input_values: Optional[torch.FloatTensor] = None,1013 attention_mask: Optional[torch.LongTensor] = None,1014 decoder_input_ids: Optional[torch.LongTensor] = None,1015 decoder_attention_mask: Optional[torch.LongTensor] = None,1016 encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,1017 past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None,1018 decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None,1019 decoder_position_ids: Optional[tuple[torch.LongTensor]] = None,1020 use_cache: Optional[bool] = None,1021 cache_position: Optional[torch.LongTensor] = None,1022 labels: Optional[torch.LongTensor] = None,1023 **kwargs: Unpack[TransformersKwargs],1024 ) -> Seq2SeqLMOutput:1025 r"""1026 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):1027 Float values of the raw speech waveform. Raw speech waveform can be1028 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a1029 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or1030 the soundfile library (`pip install soundfile`). To prepare the array into1031 `input_values`, the [`AutoFeatureExtractor`] should be used for padding1032 and conversion into a tensor of type `torch.FloatTensor`.1033 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):1034 Indices of positions of each input sequence tokens in the position embeddings.1035 Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`1036 1037 Example:1038 1039 ```python1040 >>> import torch1041 >>> from transformers import AutoProcessor, MoonshineForConditionalGeneration1042 >>> from datasets import load_dataset1043 1044 >>> processor = AutoProcessor.from_pretrained("UsefulSensors/moonshine-tiny")1045 >>> model = MoonshineForConditionalGeneration.from_pretrained("UsefulSensors/moonshine-tiny")1046 1047 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")1048 1049 >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")1050 >>> input_values = inputs.input_values1051 1052 >>> generated_ids = model.generate(input_values, max_new_tokens=100)1053 1054 >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]1055 >>> transcription1056 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'1057 ```"""1058 1059 if labels is not None:1060 if decoder_input_ids is None and decoder_inputs_embeds is None:1061 decoder_input_ids = shift_tokens_right(1062 labels, self.config.pad_token_id, self.config.decoder_start_token_id1063 )1064 1065 outputs: Seq2SeqModelOutput = self.model(1066 input_values,1067 attention_mask=attention_mask,1068 decoder_input_ids=decoder_input_ids,1069 encoder_outputs=encoder_outputs,1070 decoder_attention_mask=decoder_attention_mask,1071 past_key_values=past_key_values,1072 decoder_inputs_embeds=decoder_inputs_embeds,1073 decoder_position_ids=decoder_position_ids,1074 use_cache=use_cache,1075 cache_position=cache_position,1076 **kwargs,1077 )1078 logits = self.proj_out(outputs.last_hidden_state)1079 1080 loss = None1081 if labels is not None:1082 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)1083 1084 return Seq2SeqLMOutput(1085 loss=loss,1086 logits=logits,1087 past_key_values=outputs.past_key_values,1088 decoder_hidden_states=outputs.decoder_hidden_states,1089 decoder_attentions=outputs.decoder_attentions,1090 cross_attentions=outputs.cross_attentions,1091 encoder_last_hidden_state=outputs.encoder_last_hidden_state,1092 encoder_hidden_states=outputs.encoder_hidden_states,1093 encoder_attentions=outputs.encoder_attentions,1094 )1095 1096 1097__all__ = ["MoonshineModel", "MoonshinePreTrainedModel", "MoonshineForConditionalGeneration"]1098