Aluode/PerceptionLabPortable
0
1# Copyright 2025 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from typing import Callable, Optional, Union16 17import torch18import torch.nn as nn19 20from transformers.utils.generic import OutputRecorder, check_model_inputs21 22from ...activations import ACT2FN23from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache24from ...configuration_utils import PretrainedConfig25from ...generation import GenerationMixin26from ...masking_utils import create_causal_mask27from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa28from ...modeling_flash_attention_utils import FlashAttentionKwargs29from ...modeling_layers import GradientCheckpointingLayer30from ...modeling_outputs import (31 BaseModelOutput,32 BaseModelOutputWithPast,33 BaseModelOutputWithPastAndCrossAttentions,34 Seq2SeqLMOutput,35 Seq2SeqModelOutput,36)37from ...modeling_rope_utils import rope_config_validation38from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel39from ...processing_utils import Unpack40from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging41from ...utils.deprecation import deprecate_kwarg42from ..glm.modeling_glm import GlmAttention, GlmRotaryEmbedding, apply_rotary_pos_emb43from ..llama.modeling_llama import LlamaDecoderLayer, LlamaModel, eager_attention_forward44from ..whisper.modeling_whisper import WhisperModel, shift_tokens_right45 46 47logger = logging.get_logger(__name__)48 49 50class MoonshineConfig(PretrainedConfig):51 r"""52 This is the configuration class to store the configuration of a [`MoonshineModel`]. It is used to instantiate a Moonshine53 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the54 defaults will yield a similar configuration to that of the Moonshine55 [UsefulSensors/moonshine-tiny](https://huggingface.co/UsefulSensors/moonshine-tiny).56 57 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the58 documentation from [`PretrainedConfig`] for more information.59 60 Args:61 vocab_size (`int`, *optional*, defaults to 32768):62 Vocabulary size of the Moonshine model. Defines the number of different tokens that can be represented by the63 `inputs_ids` passed when calling [`MoonshineModel`].64 hidden_size (`int`, *optional*, defaults to 288):65 Dimension of the hidden representations.66 intermediate_size (`int`, *optional*, defaults to 1152):67 Dimension of the MLP representations.68 encoder_num_hidden_layers (`int`, *optional*, defaults to 6):69 Number of hidden layers in the Transformer encoder.70 decoder_num_hidden_layers (`int`, *optional*, defaults to 6):71 Number of hidden layers in the Transformer decoder.72 encoder_num_attention_heads (`int`, *optional*, defaults to 8):73 Number of attention heads for each attention layer in the Transformer encoder.74 decoder_num_attention_heads (`int`, *optional*, defaults to 8):75 Number of attention heads for each attention layer in the Transformer decoder.76 encoder_num_key_value_heads (`int`, *optional*):77 This is the number of key_value heads that should be used to implement Grouped Query Attention. If78 `encoder_num_key_value_heads=encoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if79 `encoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When80 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed81 by meanpooling all the original heads within that group. For more details, check out [this82 paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to83 `num_attention_heads`.84 decoder_num_key_value_heads (`int`, *optional*):85 This is the number of key_value heads that should be used to implement Grouped Query Attention. If86 `decoder_num_key_value_heads=decoder_num_attention_heads`, the model will use Multi Head Attention (MHA), if87 `decoder_num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When88 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed89 by meanpooling all the original heads within that group. For more details, check out [this90 paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to91 `decoder_num_attention_heads`.92 pad_head_dim_to_multiple_of (`int`, *optional*):93 Pad head dimension in encoder and decoder to the next multiple of this value. Necessary for using certain94 optimized attention implementations.95 encoder_hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):96 The non-linear activation function (function or string) in the encoder.97 decoder_hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):98 The non-linear activation function (function or string) in the decoder.99 max_position_embeddings (`int`, *optional*, defaults to 512):100 The maximum sequence length that this model might ever be used with.101 initializer_range (`float`, *optional*, defaults to 0.02):102 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.103 decoder_start_token_id (`int`, *optional*, defaults to 1):104 Corresponds to the "<|startoftranscript|>" token, which is automatically used when no `decoder_input_ids`105 are provided to the `generate` function. It is used to guide the model`s generation process depending on106 the task.107 use_cache (`bool`, *optional*, defaults to `True`):108 Whether or not the model should return the last key/values attentions (not used by all models).109 rope_theta (`float`, *optional*, defaults to 10000.0):110 The base period of the RoPE embeddings.111 rope_scaling (`Dict`, *optional*):112 Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type113 and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value114 accordingly.115 Expected contents:116 `rope_type` (`str`):117 The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',118 'llama3'], with 'default' being the original RoPE implementation.119 `factor` (`float`, *optional*):120 Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In121 most scaling types, a `factor` of x will enable the model to handle sequences of length x *122 original maximum pre-trained length.123 `original_max_position_embeddings` (`int`, *optional*):124 Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during125 pretraining.126 `attention_factor` (`float`, *optional*):127 Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention128 computation. If unspecified, it defaults to value recommended by the implementation, using the129 `factor` field to infer the suggested value.130 `beta_fast` (`float`, *optional*):131 Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear132 ramp function. If unspecified, it defaults to 32.133 `beta_slow` (`float`, *optional*):134 Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear135 ramp function. If unspecified, it defaults to 1.136 `short_factor` (`list[float]`, *optional*):137 Only used with 'longrope'. The scaling factor to be applied to short contexts (<138 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden139 size divided by the number of attention heads divided by 2140 `long_factor` (`list[float]`, *optional*):141 Only used with 'longrope'. The scaling factor to be applied to long contexts (<142 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden143 size divided by the number of attention heads divided by 2144 `low_freq_factor` (`float`, *optional*):145 Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE146 `high_freq_factor` (`float`, *optional*):147 Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE148 partial_rotary_factor (`float`, *optional*, defaults to 0.9):149 Percentage of the query and keys which will have rotary embedding.150 is_encoder_decoder (`bool`, *optional*, defaults to `True`):151 Whether the model is used as an encoder/decoder or not.152 attention_bias (`bool`, *optional*, defaults to `False`):153 Whether to use a bias in the query, key, value and output projection layers during self-attention.154 attention_dropout (`float`, *optional*, defaults to 0.0):155 The dropout ratio for the attention probabilities.156 bos_token_id (`int`, *optional*, defaults to 1):157 Denotes beginning of sequences token id.158 eos_token_id (`int`, *optional*, defaults to 2):159 Denotes end of sequences token id.160 161 Example:162 163 ```python164 >>> from transformers import MoonshineModel, MoonshineConfig165 166 >>> # Initializing a Moonshine style configuration167 >>> configuration = MoonshineConfig().from_pretrained("UsefulSensors/moonshine-tiny")168 169 >>> # Initializing a model from the configuration170 >>> model = MoonshineModel(configuration)171 172 >>> # Accessing the model configuration173 >>> configuration = model.config174 ```"""175 176 model_type = "moonshine"177 keys_to_ignore_at_inference = ["past_key_values"]178 attribute_map = {179 "num_key_value_heads": "encoder_num_key_value_heads",180 "num_attention_heads": "encoder_num_attention_heads",181 "num_hidden_layers": "encoder_num_hidden_layers",182 }183 184 def __init__(185 self,186 vocab_size=32768,187 hidden_size=288,188 intermediate_size=1152,189 encoder_num_hidden_layers=6,190 decoder_num_hidden_layers=6,191 encoder_num_attention_heads=8,192 decoder_num_attention_heads=8,193 encoder_num_key_value_heads=None,194 decoder_num_key_value_heads=None,195 pad_head_dim_to_multiple_of=None,196 encoder_hidden_act="gelu",197 decoder_hidden_act="silu",198 max_position_embeddings=512,199 initializer_range=0.02,200 decoder_start_token_id=1,201 use_cache=True,202 rope_theta=10000.0,203 rope_scaling=None,204 partial_rotary_factor=0.9,205 is_encoder_decoder=True,206 attention_bias=False,207 attention_dropout=0.0,208 bos_token_id=1,209 eos_token_id=2,210 **kwargs,211 ):212 self.vocab_size = vocab_size213 self.hidden_size = hidden_size214 self.intermediate_size = intermediate_size215 self.encoder_num_hidden_layers = encoder_num_hidden_layers216 self.decoder_num_hidden_layers = decoder_num_hidden_layers217 self.encoder_num_attention_heads = encoder_num_attention_heads218 self.decoder_num_attention_heads = decoder_num_attention_heads219 220 if encoder_num_key_value_heads is None:221 encoder_num_key_value_heads = encoder_num_attention_heads222 self.encoder_num_key_value_heads = encoder_num_key_value_heads223 224 if decoder_num_key_value_heads is None:225 decoder_num_key_value_heads = decoder_num_attention_heads226 self.decoder_num_key_value_heads = decoder_num_key_value_heads227 228 self.pad_head_dim_to_multiple_of = pad_head_dim_to_multiple_of229 230 self.encoder_hidden_act = encoder_hidden_act231 self.decoder_hidden_act = decoder_hidden_act232 self.max_position_embeddings = max_position_embeddings233 self.initializer_range = initializer_range234 self.decoder_start_token_id = decoder_start_token_id235 self.use_cache = use_cache236 self.rope_theta = rope_theta237 self.rope_scaling = rope_scaling238 self.partial_rotary_factor = partial_rotary_factor239 self.is_encoder_decoder = is_encoder_decoder240 self.attention_bias = attention_bias241 self.attention_dropout = attention_dropout242 243 # Validate the correctness of rotary position embeddings parameters244 rope_config_validation(self)245 246 super().__init__(247 bos_token_id=bos_token_id,248 eos_token_id=eos_token_id,249 is_encoder_decoder=is_encoder_decoder,250 decoder_start_token_id=decoder_start_token_id,251 **kwargs,252 )253 254 255class MoonshineEncoderMLP(nn.Module):256 def __init__(self, config, hidden_act):257 super().__init__()258 self.config = config259 self.activation_fn = ACT2FN[hidden_act]260 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)261 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)262 263 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:264 hidden_states = self.fc1(hidden_states)265 hidden_states = self.activation_fn(hidden_states)266 hidden_states = self.fc2(hidden_states)267 return hidden_states268 269 270class MoonshineDecoderMLP(nn.Module):271 def __init__(self, config, hidden_act):272 super().__init__()273 self.config = config274 self.activation_fn = ACT2FN[hidden_act]275 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size * 2)276 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)277 278 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:279 hidden_states = self.fc1(hidden_states)280 hidden_states, gate = hidden_states.chunk(2, dim=-1)281 hidden_states = self.activation_fn(gate) * hidden_states282 hidden_states = self.fc2(hidden_states)283 return hidden_states284 285 286class MoonshineAttention(GlmAttention):287 def __init__(288 self,289 config: MoonshineConfig,290 layer_idx: int,291 is_causal: bool,292 num_attention_heads: int,293 num_key_value_heads: int,294 ):295 config.update({"num_attention_heads": num_attention_heads, "num_key_value_heads": num_key_value_heads})296 super().__init__(config, layer_idx)297 self.is_causal = is_causal298 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)299 300 # Pad head dimension to the next specified multiple.301 if self.config.pad_head_dim_to_multiple_of is not None:302 target_multiple = self.config.pad_head_dim_to_multiple_of303 target_head_dim = target_multiple * ((self.head_dim + target_multiple - 1) // target_multiple)304 self.head_dim_padding = target_head_dim - self.head_dim305 else:306 self.head_dim_padding = 0307 308 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")309 def forward(310 self,311 hidden_states: torch.Tensor,312 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,313 attention_mask: Optional[torch.Tensor] = None,314 past_key_values: Optional[Cache] = None,315 cache_position: Optional[torch.LongTensor] = None,316 key_value_states: Optional[torch.Tensor] = None,317 **kwargs: Unpack[FlashAttentionKwargs],318 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:319 bsz, q_len = hidden_states.shape[:-1]320 321 query_states = (322 self.q_proj(hidden_states).view(bsz, q_len, self.config.num_key_value_heads, self.head_dim).transpose(1, 2)323 )324 325 is_cross_attention = key_value_states is not None326 if past_key_values is not None:327 is_updated = past_key_values.is_updated.get(self.layer_idx)328 if is_cross_attention:329 # after the first generated id, we can subsequently re-use all key/value_states from cache330 past_key_values.is_updated[self.layer_idx] = True331 past_key_values = past_key_values.cross_attention_cache332 else:333 past_key_values = past_key_values.self_attention_cache334 335 # use key_value_states if cross attention336 current_states = key_value_states if key_value_states is not None else hidden_states337 if is_cross_attention and past_key_values and is_updated:338 key_states = past_key_values.layers[self.layer_idx].keys339 value_states = past_key_values.layers[self.layer_idx].values340 else:341 key_states = (342 self.k_proj(current_states)343 .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)344 .transpose(1, 2)345 )346 value_states = (347 self.v_proj(current_states)348 .view(bsz, -1, self.config.num_key_value_heads, self.head_dim)349 .transpose(1, 2)350 )351 if is_cross_attention and past_key_values is not None:352 key_states, value_states = past_key_values.update(353 key_states, value_states, self.layer_idx, {"cache_position": cache_position}354 )355 356 if not is_cross_attention:357 cos, sin = position_embeddings358 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)359 360 if past_key_values is not None:361 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}362 key_states, value_states = past_key_values.update(363 key_states, value_states, self.layer_idx, cache_kwargs364 )365 366 attention_interface: Callable = eager_attention_forward367 if self.config._attn_implementation != "eager":368 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]369 370 is_causal = self.is_causal and attention_mask is None and q_len > 1371 372 if self.head_dim_padding > 0:373 query_states = torch.nn.functional.pad(query_states, (0, self.head_dim_padding))374 key_states = torch.nn.functional.pad(key_states, (0, self.head_dim_padding))375 value_states = torch.nn.functional.pad(value_states, (0, self.head_dim_padding))376 377 attn_output, attn_weights = attention_interface(378 self,379 query_states,380 key_states,381 value_states,382 attention_mask,383 dropout=0.0 if not self.training else self.attention_dropout,384 scaling=self.scaling,385 is_causal=is_causal,386 **kwargs,387 )388 389 if self.head_dim_padding > 0:390 attn_output = attn_output[..., : -self.head_dim_padding]391 392 attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()393 attn_output = self.o_proj(attn_output)394 return attn_output, attn_weights395 396 397class MoonshineRotaryEmbedding(GlmRotaryEmbedding):398 pass399 400 401class MoonshineEncoderLayer(LlamaDecoderLayer):402 def __init__(self, config: MoonshineConfig, layer_idx: int):403 super().__init__(config, layer_idx)404 405 self.self_attn = MoonshineAttention(406 config=config,407 layer_idx=layer_idx,408 is_causal=False,409 num_attention_heads=config.encoder_num_attention_heads,410 num_key_value_heads=config.encoder_num_key_value_heads,411 )412 413 self.mlp = MoonshineEncoderMLP(config, config.encoder_hidden_act)414 self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)415 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)416 417 418class MoonshineDecoderLayer(GradientCheckpointingLayer):419 def __init__(self, config: MoonshineConfig, layer_idx: Optional[int] = None):420 super().__init__()421 self.hidden_size = config.hidden_size422 423 self.self_attn = MoonshineAttention(424 config=config,425 layer_idx=layer_idx,426 is_causal=True,427 num_attention_heads=config.decoder_num_attention_heads,428 num_key_value_heads=config.decoder_num_key_value_heads,429 )430 self.encoder_attn = MoonshineAttention(431 config=config,432 layer_idx=layer_idx,433 is_causal=False,434 num_attention_heads=config.decoder_num_attention_heads,435 num_key_value_heads=config.decoder_num_key_value_heads,436 )437 438 self.mlp = MoonshineDecoderMLP(config, config.decoder_hidden_act)439 self.input_layernorm = nn.LayerNorm(config.hidden_size, bias=False)440 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, bias=False)441 self.final_layernorm = nn.LayerNorm(config.hidden_size, bias=False)442 443 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")444 def forward(445 self,446 hidden_states: torch.Tensor,447 attention_mask: Optional[torch.Tensor] = None,448 encoder_hidden_states: Optional[torch.Tensor] = None,449 encoder_attention_mask: Optional[torch.Tensor] = None,450 position_ids: Optional[torch.LongTensor] = None,451 encoder_position_ids: Optional[torch.LongTensor] = None,452 past_key_values: Optional[Cache] = None,453 use_cache: Optional[bool] = False,454 cache_position: Optional[torch.LongTensor] = None,455 position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,456 encoder_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,457 **kwargs: Unpack[TransformersKwargs],458 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:459 residual = hidden_states460 hidden_states = self.input_layernorm(hidden_states)461 462 hidden_states, _ = self.self_attn(463 hidden_states=hidden_states,464 attention_mask=attention_mask,465 position_ids=position_ids,466 past_key_values=past_key_values,467 use_cache=use_cache,468 cache_position=cache_position,469 position_embeddings=position_embeddings,470 **kwargs,471 )472 hidden_states = residual + hidden_states473 474 if encoder_hidden_states is not None:475 residual = hidden_states476 hidden_states = self.post_attention_layernorm(hidden_states)477 hidden_states, _ = self.encoder_attn(478 hidden_states=hidden_states,479 key_value_states=encoder_hidden_states,480 attention_mask=encoder_attention_mask,481 past_key_values=past_key_values,482 use_cache=use_cache,483 )484 hidden_states = residual + hidden_states485 486 residual = hidden_states487 hidden_states = self.final_layernorm(hidden_states)488 hidden_states = self.mlp(hidden_states)489 hidden_states = residual + hidden_states490 return hidden_states491 492 493@auto_docstring494class MoonshinePreTrainedModel(PreTrainedModel):495 config: MoonshineConfig496 base_model_prefix = "model"497 main_input_name = "input_values"498 supports_gradient_checkpointing = True499 _no_split_modules = ["MoonshineEncoderLayer", "MoonshineDecoderLayer"]500 _supports_flash_attn = True501 _supports_sdpa = True502 503 _can_compile_fullgraph = True504 # TODO arthur, how do we separate when it cross / self coming from different layer?505 506 def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):507 """508 Computes the output length of the convolutional layers509 """510 output_conv1_length = int((input_lengths - 127) / 64 + 1)511 output_conv2_length = int((output_conv1_length - 7) / 3 + 1)512 output_conv3_length = int((output_conv2_length - 3) / 2 + 1)513 514 return output_conv3_length515 516 517class MoonshineEncoder(MoonshinePreTrainedModel):518 """519 Transformer encoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MoonshineEncoderLayer`]520 521 Args:522 config: MoonshineConfig523 """524 525 main_input_name = "input_values"526 _can_record_outputs = {527 "attentions": MoonshineAttention,528 "hidden_states": MoonshineEncoderLayer,529 }530 531 def __init__(self, config: MoonshineConfig):532 super().__init__(config)533 self.config = config534 embed_dim = config.hidden_size535 536 self.conv1 = nn.Conv1d(1, embed_dim, kernel_size=127, stride=64, bias=False)537 self.conv2 = nn.Conv1d(embed_dim, 2 * embed_dim, kernel_size=7, stride=3)538 self.conv3 = nn.Conv1d(2 * embed_dim, embed_dim, kernel_size=3, stride=2)539 self.groupnorm = nn.GroupNorm(num_groups=1, num_channels=embed_dim, eps=1e-5)540 self.rotary_emb = MoonshineRotaryEmbedding(config=config)541 542 self.layers = nn.ModuleList(543 [MoonshineEncoderLayer(config, idx) for idx in range(config.encoder_num_hidden_layers)]544 )545 self.layer_norm = nn.LayerNorm(embed_dim, bias=False)546 self.gradient_checkpointing = False547 self.post_init()548 549 def get_input_embeddings(self) -> nn.Module:550 return self.conv1551 552 def set_input_embeddings(self, value: nn.Module):553 self.conv1 = value554 555 @check_model_inputs()556 def forward(557 self,558 input_values: torch.FloatTensor,559 attention_mask: Optional[torch.Tensor] = None,560 **kwargs: Unpack[TransformersKwargs],561 ) -> BaseModelOutputWithPast:562 r"""563 Args:564 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):565 Float values of the raw speech waveform. Raw speech waveform can be566 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a567 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or568 the soundfile library (`pip install soundfile`). To prepare the array into569 `input_values`, the [`AutoFeatureExtractor`] should be used for padding570 and conversion into a tensor of type `torch.FloatTensor`.571 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):572 Mask to avoid performing attention on padding indices in `input_values`. Mask values selected in `[0, 1]`:573 - 1 for tokens that are **not masked**,574 - 0 for tokens that are **masked**.575 [What are attention masks?](../glossary#attention-mask)576 """577 input_values = input_values.unsqueeze(1)578 hidden_states = nn.functional.tanh(self.conv1(input_values))579 hidden_states = self.groupnorm(hidden_states)580 hidden_states = nn.functional.gelu(self.conv2(hidden_states))581 hidden_states = nn.functional.gelu(self.conv3(hidden_states))582 hidden_states = hidden_states.permute(0, 2, 1)583 584 # attention mask downsampling585 if attention_mask is not None:586 mask_len = self._get_feat_extract_output_lengths(attention_mask.shape[-1])587 downsample_stride = 64 * 3 * 2 # conv strides588 attention_mask = attention_mask[..., ::downsample_stride][..., :mask_len]589 if self.config._attn_implementation == "flash_attention_2":590 attention_mask = attention_mask if (attention_mask == 0.0).any() else None591 elif self.config._attn_implementation == "sdpa":592 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, hidden_states.dtype)593 else:594 attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)595 596 position_ids = torch.arange(0, hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)597 position_embeddings = self.rotary_emb(hidden_states, position_ids)598 599 for encoder_layer in self.layers:600 hidden_states = encoder_layer(601 hidden_states,602 attention_mask=attention_mask,603 position_ids=position_ids,604 position_embeddings=position_embeddings,605 **kwargs,606 )607 608 hidden_states = self.layer_norm(hidden_states)609 610 return BaseModelOutputWithPast(611 last_hidden_state=hidden_states,612 )613 614 615class MoonshineDecoder(LlamaModel):616 main_input_name = "input_ids"617 _can_record_outputs = {618 "attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="self_attn"),619 "hidden_states": MoonshineDecoderLayer,620 "cross_attentions": OutputRecorder(MoonshineAttention, index=1, layer_name="encoder_attn"),621 }622 623 def __init__(self, config: MoonshineConfig):624 super().__init__(config)625 self.norm = nn.LayerNorm(config.hidden_size, bias=False)626 self.layers = nn.ModuleList(627 [MoonshineDecoderLayer(config, idx) for idx in range(config.decoder_num_hidden_layers)]628 )629 630 @check_model_inputs()631 def forward(632 self,633 input_ids: Optional[torch.LongTensor] = None,634 attention_mask: Optional[torch.Tensor] = None,635 position_ids: Optional[torch.LongTensor] = None,636 past_key_values: Optional[Cache] = None,637 inputs_embeds: Optional[torch.FloatTensor] = None,638 use_cache: Optional[bool] = None,639 cache_position: Optional[torch.LongTensor] = None,640 encoder_hidden_states: Optional[torch.FloatTensor] = None,641 encoder_attention_mask: Optional[torch.Tensor] = None,642 **kwargs: Unpack[TransformersKwargs],643 ) -> Union[tuple, BaseModelOutputWithPast]:644 r"""645 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):646 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention647 of the decoder.648 encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):649 Mask to avoid performing attention on padding indices in `encoder_hidden_states`. Mask values selected in `[0, 1]`:650 - 1 for tokens that are **not masked**,651 - 0 for tokens that are **masked**.652 [What are attention masks?](../glossary#attention-mask)653 """654 if (input_ids is None) ^ (inputs_embeds is not None):655 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")656 657 if inputs_embeds is None:658 inputs_embeds = self.embed_tokens(input_ids)659 660 if use_cache and past_key_values is None:661 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))662 663 if cache_position is None:664 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0665 cache_position = torch.arange(666 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device667 )668 669 if position_ids is None:670 position_ids = cache_position.unsqueeze(0)671 672 causal_mask = create_causal_mask(673 config=self.config,674 input_embeds=inputs_embeds,675 attention_mask=attention_mask,676 cache_position=cache_position,677 past_key_values=past_key_values,678 position_ids=position_ids,679 )680 681 hidden_states = inputs_embeds682 position_embeddings = self.rotary_emb(hidden_states, position_ids)683 684 if encoder_attention_mask is not None:685 mask_len = encoder_hidden_states.shape[-2]686 downsample_stride = 64 * 3 * 2 # conv strides687 encoder_attention_mask = encoder_attention_mask[..., ::downsample_stride][..., :mask_len]688 if self.config._attn_implementation == "flash_attention_2":689 encoder_attention_mask = encoder_attention_mask if (encoder_attention_mask == 0.0).any() else None690 elif self.config._attn_implementation == "sdpa":691 encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(692 encoder_attention_mask, hidden_states.dtype, hidden_states.shape[-2]693 )694 else:695 encoder_attention_mask = _prepare_4d_attention_mask(696 encoder_attention_mask, hidden_states.dtype, hidden_states.shape[-2]697 )698 699 for decoder_layer in self.layers:700 hidden_states = decoder_layer(701 hidden_states,702 causal_mask,703 encoder_hidden_states, # as a positional argument for gradient checkpointing704 encoder_attention_mask=encoder_attention_mask,705 position_ids=position_ids,706 past_key_values=past_key_values,707 use_cache=use_cache,708 cache_position=cache_position,709 position_embeddings=position_embeddings,710 **kwargs,711 )712 713 hidden_states = self.norm(hidden_states)714 715 return BaseModelOutputWithPastAndCrossAttentions(716 last_hidden_state=hidden_states,717 past_key_values=past_key_values if use_cache else None,718 )719 720 721class MoonshineModel(WhisperModel):722 @can_return_tuple723 @auto_docstring724 def forward(725 self,726 input_values: Optional[torch.FloatTensor] = None,727 attention_mask: Optional[torch.LongTensor] = None,728 decoder_input_ids: Optional[torch.LongTensor] = None,729 decoder_attention_mask: Optional[torch.LongTensor] = None,730 encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,731 past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None,732 decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None,733 decoder_position_ids: Optional[tuple[torch.LongTensor]] = None,734 use_cache: Optional[bool] = None,735 cache_position: Optional[torch.LongTensor] = None,736 **kwargs: Unpack[TransformersKwargs],737 ) -> Seq2SeqModelOutput:738 r"""739 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):740 Float values of the raw speech waveform. Raw speech waveform can be741 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a742 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or743 the soundfile library (`pip install soundfile`). To prepare the array into744 `input_values`, the [`AutoFeatureExtractor`] should be used for padding745 and conversion into a tensor of type `torch.FloatTensor`.746 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):747 Indices of positions of each input sequence tokens in the position embeddings.748 Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`749 750 Example:751 752 ```python753 >>> import torch754 >>> from transformers import AutoFeatureExtractor, MoonshineModel755 >>> from datasets import load_dataset756 757 >>> model = MoonshineModel.from_pretrained("UsefulSensors/moonshine-tiny")758 >>> feature_extractor = AutoFeatureExtractor.from_pretrained("UsefulSensors/moonshine-tiny")759 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")760 >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")761 >>> input_values = inputs.input_values762 >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id763 >>> last_hidden_state = model(input_values, decoder_input_ids=decoder_input_ids).last_hidden_state764 >>> list(last_hidden_state.shape)765 [1, 2, 288]766 ```767 """768 if encoder_outputs is None:769 encoder_outputs: BaseModelOutput = self.encoder(input_values, attention_mask=attention_mask, **kwargs)770 771 decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(772 input_ids=decoder_input_ids,773 attention_mask=decoder_attention_mask,774 encoder_attention_mask=attention_mask,775 encoder_hidden_states=encoder_outputs.last_hidden_state,776 past_key_values=past_key_values,777 inputs_embeds=decoder_inputs_embeds,778 position_ids=decoder_position_ids,779 use_cache=use_cache,780 cache_position=cache_position,781 **kwargs,782 )783 784 return Seq2SeqModelOutput(785 last_hidden_state=decoder_outputs.last_hidden_state,786 past_key_values=decoder_outputs.past_key_values,787 decoder_hidden_states=decoder_outputs.hidden_states,788 decoder_attentions=decoder_outputs.attentions,789 cross_attentions=decoder_outputs.cross_attentions,790 encoder_last_hidden_state=encoder_outputs.last_hidden_state,791 encoder_hidden_states=encoder_outputs.hidden_states,792 encoder_attentions=encoder_outputs.attentions,793 )794 795 796@auto_docstring(797 custom_intro="""798 The Moonshine Model with a language modeling head. Can be used for automatic speech recognition.799 """800)801class MoonshineForConditionalGeneration(MoonshinePreTrainedModel, GenerationMixin):802 _tied_weights_keys = ["proj_out.weight"]803 804 def __init__(self, config: MoonshineConfig):805 super().__init__(config)806 self.model = MoonshineModel(config)807 self.proj_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)808 809 # Initialize weights and apply final processing810 self.post_init()811 812 def get_encoder(self):813 return self.model.get_encoder()814 815 def get_decoder(self):816 return self.model.get_decoder()817 818 def get_output_embeddings(self):819 return self.proj_out820 821 def set_output_embeddings(self, new_embeddings):822 self.proj_out = new_embeddings823 824 def get_input_embeddings(self) -> nn.Module:825 return self.model.get_input_embeddings()826 827 @can_return_tuple828 @auto_docstring829 def forward(830 self,831 input_values: Optional[torch.FloatTensor] = None,832 attention_mask: Optional[torch.LongTensor] = None,833 decoder_input_ids: Optional[torch.LongTensor] = None,834 decoder_attention_mask: Optional[torch.LongTensor] = None,835 encoder_outputs: Optional[tuple[tuple[torch.FloatTensor]]] = None,836 past_key_values: Optional[Union[EncoderDecoderCache, tuple[torch.FloatTensor]]] = None,837 decoder_inputs_embeds: Optional[tuple[torch.FloatTensor]] = None,838 decoder_position_ids: Optional[tuple[torch.LongTensor]] = None,839 use_cache: Optional[bool] = None,840 cache_position: Optional[torch.LongTensor] = None,841 labels: Optional[torch.LongTensor] = None,842 **kwargs: Unpack[TransformersKwargs],843 ) -> Seq2SeqLMOutput:844 r"""845 input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`):846 Float values of the raw speech waveform. Raw speech waveform can be847 obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]`, a848 `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or849 the soundfile library (`pip install soundfile`). To prepare the array into850 `input_values`, the [`AutoFeatureExtractor`] should be used for padding851 and conversion into a tensor of type `torch.FloatTensor`.852 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`):853 Indices of positions of each input sequence tokens in the position embeddings.854 Used to calculate the position embeddings up to `config.decoder_config.max_position_embeddings`855 856 Example:857 858 ```python859 >>> import torch860 >>> from transformers import AutoProcessor, MoonshineForConditionalGeneration861 >>> from datasets import load_dataset862 863 >>> processor = AutoProcessor.from_pretrained("UsefulSensors/moonshine-tiny")864 >>> model = MoonshineForConditionalGeneration.from_pretrained("UsefulSensors/moonshine-tiny")865 866 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")867 868 >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt")869 >>> input_values = inputs.input_values870 871 >>> generated_ids = model.generate(input_values, max_new_tokens=100)872 873 >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]874 >>> transcription875 'Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'876 ```"""877 878 if labels is not None:879 if decoder_input_ids is None and decoder_inputs_embeds is None:880 decoder_input_ids = shift_tokens_right(881 labels, self.config.pad_token_id, self.config.decoder_start_token_id882 )883 884 outputs: Seq2SeqModelOutput = self.model(885 input_values,886 attention_mask=attention_mask,887 decoder_input_ids=decoder_input_ids,888 encoder_outputs=encoder_outputs,889 decoder_attention_mask=decoder_attention_mask,890 past_key_values=past_key_values,891 decoder_inputs_embeds=decoder_inputs_embeds,892 decoder_position_ids=decoder_position_ids,893 use_cache=use_cache,894 cache_position=cache_position,895 **kwargs,896 )897 logits = self.proj_out(outputs.last_hidden_state)898 899 loss = None900 if labels is not None:901 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)902 903 return Seq2SeqLMOutput(904 loss=loss,905 logits=logits,906 past_key_values=outputs.past_key_values,907 decoder_hidden_states=outputs.decoder_hidden_states,908 decoder_attentions=outputs.decoder_attentions,909 cross_attentions=outputs.cross_attentions,910 encoder_last_hidden_state=outputs.encoder_last_hidden_state,911 encoder_hidden_states=outputs.encoder_hidden_states,912 encoder_attentions=outputs.encoder_attentions,913 )914 915 916__all__ = [917 "MoonshineConfig",918 "MoonshineModel",919 "MoonshinePreTrainedModel",920 "MoonshineForConditionalGeneration",921]922