Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch Wav2Vec2 model."""16 17import math18import warnings19from dataclasses import dataclass20from typing import Callable, Optional, Union21 22import numpy as np23import torch24from safetensors.torch import load_file as safe_load_file25from torch import nn26from torch.nn import CrossEntropyLoss27 28from ...activations import ACT2FN29from ...integrations.deepspeed import is_deepspeed_zero3_enabled30from ...integrations.fsdp import is_fsdp_managed_module31from ...modeling_attn_mask_utils import (32 _prepare_4d_attention_mask,33 _prepare_4d_attention_mask_for_sdpa,34)35from ...modeling_flash_attention_utils import FlashAttentionKwargs36from ...modeling_layers import GradientCheckpointingLayer37from ...modeling_outputs import (38 BaseModelOutput,39 CausalLMOutput,40 MaskedLMOutput,41 SequenceClassifierOutput,42 TokenClassifierOutput,43 Wav2Vec2BaseModelOutput,44 XVectorOutput,45)46from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel47from ...processing_utils import Unpack48from ...utils import (49 ModelOutput,50 auto_docstring,51 cached_file,52 check_torch_load_is_safe,53 is_peft_available,54 is_torch_flex_attn_available,55 logging,56)57from .configuration_wav2vec2 import Wav2Vec2Config58 59 60WAV2VEC2_ADAPTER_PT_FILE = "adapter.{}.bin"61WAV2VEC2_ADAPTER_SAFE_FILE = "adapter.{}.safetensors"62 63if is_torch_flex_attn_available():64 from ...integrations.flex_attention import make_flex_block_causal_mask65 66 67logger = logging.get_logger(__name__)68 69 70_HIDDEN_STATES_START_POSITION = 271 72 73@dataclass74@auto_docstring(75 custom_intro="""76 Output type of [`Wav2Vec2ForPreTraining`], with potential hidden states and attentions.77 """78)79class Wav2Vec2ForPreTrainingOutput(ModelOutput):80 r"""81 loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):82 Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official83 paper](https://huggingface.co/papers/2006.11477).84 projected_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):85 Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked86 projected quantized states.87 projected_quantized_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):88 Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive89 target vectors for contrastive loss.90 codevector_perplexity (`torch.FloatTensor` of shape `(1,)`):91 The perplexity of the codevector distribution, used to measure the diversity of the codebook.92 contrastive_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):93 The contrastive loss (L_m) as stated in the [official paper](https://huggingface.co/papers/2006.11477).94 diversity_loss (*optional*, returned when `sample_negative_indices` are passed, `torch.FloatTensor` of shape `(1,)`):95 The diversity loss (L_d) as stated in the [official paper](https://huggingface.co/papers/2006.11477).96 """97 98 loss: Optional[torch.FloatTensor] = None99 projected_states: Optional[torch.FloatTensor] = None100 projected_quantized_states: Optional[torch.FloatTensor] = None101 codevector_perplexity: Optional[torch.FloatTensor] = None102 hidden_states: Optional[tuple[torch.FloatTensor]] = None103 attentions: Optional[tuple[torch.FloatTensor]] = None104 contrastive_loss: Optional[torch.FloatTensor] = None105 diversity_loss: Optional[torch.FloatTensor] = None106 107 108def _compute_mask_indices(109 shape: tuple[int, int],110 mask_prob: float,111 mask_length: int,112 attention_mask: Optional[torch.LongTensor] = None,113 min_masks: int = 0,114) -> np.ndarray:115 """116 Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for117 ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on118 CPU as part of the preprocessing during training.119 120 Args:121 shape: The shape for which to compute masks. This should be of a tuple of size 2 where122 the first element is the batch size and the second element is the length of the axis to span.123 mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of124 independently generated mask spans of length `mask_length` is computed by125 `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the126 actual percentage will be smaller.127 mask_length: size of the mask128 min_masks: minimum number of masked spans129 attention_mask: A (right-padded) attention mask which independently shortens the feature axis of130 each batch dimension.131 """132 batch_size, sequence_length = shape133 134 if mask_length < 1:135 raise ValueError("`mask_length` has to be bigger than 0.")136 137 if mask_length > sequence_length:138 raise ValueError(139 f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"140 f" and `sequence_length`: {sequence_length}`"141 )142 143 # epsilon is used for probabilistic rounding144 epsilon = np.random.rand(1).item()145 146 def compute_num_masked_span(input_length):147 """Given input length, compute how many spans should be masked"""148 num_masked_span = int(mask_prob * input_length / mask_length + epsilon)149 num_masked_span = max(num_masked_span, min_masks)150 151 # make sure num masked span <= sequence_length152 if num_masked_span * mask_length > sequence_length:153 num_masked_span = sequence_length // mask_length154 155 # make sure num_masked span is also <= input_length - (mask_length - 1)156 if input_length - (mask_length - 1) < num_masked_span:157 num_masked_span = max(input_length - (mask_length - 1), 0)158 159 return num_masked_span160 161 # compute number of masked spans in batch162 input_lengths = (163 attention_mask.detach().sum(-1).tolist()164 if attention_mask is not None165 else [sequence_length for _ in range(batch_size)]166 )167 168 # SpecAugment mask to fill169 spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)170 spec_aug_mask_idxs = []171 172 max_num_masked_span = compute_num_masked_span(sequence_length)173 174 if max_num_masked_span == 0:175 return spec_aug_mask176 177 for input_length in input_lengths:178 # compute num of masked spans for this input179 num_masked_span = compute_num_masked_span(input_length)180 181 # get random indices to mask182 spec_aug_mask_idx = np.random.choice(183 np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False184 )185 186 # pick first sampled index that will serve as a dummy index to pad vector187 # to ensure same dimension for all batches due to probabilistic rounding188 # Picking first sample just pads those vectors twice.189 if len(spec_aug_mask_idx) == 0:190 # this case can only happen if `input_length` is strictly smaller then191 # `sequence_length` in which case the last token has to be a padding192 # token which we can use as a dummy mask id193 dummy_mask_idx = sequence_length - 1194 else:195 dummy_mask_idx = spec_aug_mask_idx[0]196 197 spec_aug_mask_idx = np.concatenate(198 [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]199 )200 spec_aug_mask_idxs.append(spec_aug_mask_idx)201 202 spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)203 204 # expand masked indices to masked spans205 spec_aug_mask_idxs = np.broadcast_to(206 spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)207 )208 spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)209 210 # add offset to the starting indexes so that indexes now create a span211 offsets = np.arange(mask_length)[None, None, :]212 offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(213 batch_size, max_num_masked_span * mask_length214 )215 spec_aug_mask_idxs = spec_aug_mask_idxs + offsets216 217 # ensure that we cannot have indices larger than sequence_length218 if spec_aug_mask_idxs.max() > sequence_length - 1:219 spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1220 221 # scatter indices to mask222 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)223 224 return spec_aug_mask225 226 227def _sample_negative_indices(228 features_shape: tuple, num_negatives: int, mask_time_indices: Optional[np.ndarray] = None229):230 """231 Sample `num_negatives` vectors from feature vectors.232 """233 batch_size, sequence_length = features_shape234 235 # generate indices of the positive vectors themselves, repeat them `num_negatives` times236 sequence_length_range = np.arange(sequence_length)237 238 # get `num_negatives` random vector indices from the same utterance239 sampled_negative_indices = np.zeros(shape=(batch_size, sequence_length, num_negatives), dtype=np.int32)240 241 mask_time_indices = (242 mask_time_indices.astype(bool) if mask_time_indices is not None else np.ones(features_shape, dtype=bool)243 )244 245 for batch_idx in range(batch_size):246 high = mask_time_indices[batch_idx].sum() - 1247 mapped_masked_indices = sequence_length_range[mask_time_indices[batch_idx]]248 249 feature_indices = np.broadcast_to(np.arange(high + 1)[:, None], (high + 1, num_negatives))250 sampled_indices = np.random.randint(0, high, size=(high + 1, num_negatives))251 # avoid sampling the same positive vector, but keep the distribution uniform252 sampled_indices[sampled_indices >= feature_indices] += 1253 254 # remap to actual indices255 sampled_negative_indices[batch_idx][mask_time_indices[batch_idx]] = mapped_masked_indices[sampled_indices]256 257 # correct for batch size258 sampled_negative_indices[batch_idx] += batch_idx * sequence_length259 260 return sampled_negative_indices261 262 263class Wav2Vec2NoLayerNormConvLayer(GradientCheckpointingLayer):264 def __init__(self, config, layer_id=0):265 super().__init__()266 self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1267 self.out_conv_dim = config.conv_dim[layer_id]268 269 self.conv = nn.Conv1d(270 self.in_conv_dim,271 self.out_conv_dim,272 kernel_size=config.conv_kernel[layer_id],273 stride=config.conv_stride[layer_id],274 bias=config.conv_bias,275 )276 self.activation = ACT2FN[config.feat_extract_activation]277 278 def forward(self, hidden_states):279 hidden_states = self.conv(hidden_states)280 hidden_states = self.activation(hidden_states)281 return hidden_states282 283 284class Wav2Vec2LayerNormConvLayer(GradientCheckpointingLayer):285 def __init__(self, config, layer_id=0):286 super().__init__()287 self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1288 self.out_conv_dim = config.conv_dim[layer_id]289 290 self.conv = nn.Conv1d(291 self.in_conv_dim,292 self.out_conv_dim,293 kernel_size=config.conv_kernel[layer_id],294 stride=config.conv_stride[layer_id],295 bias=config.conv_bias,296 )297 self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True)298 self.activation = ACT2FN[config.feat_extract_activation]299 300 def forward(self, hidden_states):301 hidden_states = self.conv(hidden_states)302 303 hidden_states = hidden_states.transpose(-2, -1)304 hidden_states = self.layer_norm(hidden_states)305 hidden_states = hidden_states.transpose(-2, -1)306 307 hidden_states = self.activation(hidden_states)308 return hidden_states309 310 311class Wav2Vec2GroupNormConvLayer(GradientCheckpointingLayer):312 def __init__(self, config, layer_id=0):313 super().__init__()314 self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1315 self.out_conv_dim = config.conv_dim[layer_id]316 317 self.conv = nn.Conv1d(318 self.in_conv_dim,319 self.out_conv_dim,320 kernel_size=config.conv_kernel[layer_id],321 stride=config.conv_stride[layer_id],322 bias=config.conv_bias,323 )324 self.activation = ACT2FN[config.feat_extract_activation]325 326 self.layer_norm = nn.GroupNorm(num_groups=self.out_conv_dim, num_channels=self.out_conv_dim, affine=True)327 328 def forward(self, hidden_states):329 hidden_states = self.conv(hidden_states)330 hidden_states = self.layer_norm(hidden_states)331 hidden_states = self.activation(hidden_states)332 return hidden_states333 334 335class Wav2Vec2PositionalConvEmbedding(nn.Module):336 def __init__(self, config):337 super().__init__()338 self.conv = nn.Conv1d(339 config.hidden_size,340 config.hidden_size,341 kernel_size=config.num_conv_pos_embeddings,342 padding=config.num_conv_pos_embeddings // 2,343 groups=config.num_conv_pos_embedding_groups,344 )345 346 weight_norm = nn.utils.weight_norm347 if hasattr(nn.utils.parametrizations, "weight_norm"):348 weight_norm = nn.utils.parametrizations.weight_norm349 350 if is_deepspeed_zero3_enabled():351 import deepspeed352 353 with deepspeed.zero.GatheredParameters(self.conv.weight, modifier_rank=0):354 self.conv = weight_norm(self.conv, name="weight", dim=2)355 if hasattr(self.conv, "parametrizations"):356 weight_g = self.conv.parametrizations.weight.original0357 weight_v = self.conv.parametrizations.weight.original1358 else:359 weight_g = self.conv.weight_g360 weight_v = self.conv.weight_v361 deepspeed.zero.register_external_parameter(self, weight_v)362 deepspeed.zero.register_external_parameter(self, weight_g)363 else:364 self.conv = weight_norm(self.conv, name="weight", dim=2)365 366 self.padding = Wav2Vec2SamePadLayer(config.num_conv_pos_embeddings)367 self.activation = ACT2FN[config.feat_extract_activation]368 369 def forward(self, hidden_states):370 hidden_states = hidden_states.transpose(1, 2)371 372 hidden_states = self.conv(hidden_states)373 hidden_states = self.padding(hidden_states)374 hidden_states = self.activation(hidden_states)375 376 hidden_states = hidden_states.transpose(1, 2)377 return hidden_states378 379 380class Wav2Vec2SamePadLayer(nn.Module):381 def __init__(self, num_conv_pos_embeddings):382 super().__init__()383 self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 0384 385 def forward(self, hidden_states):386 if self.num_pad_remove > 0:387 hidden_states = hidden_states[:, :, : -self.num_pad_remove]388 return hidden_states389 390 391class Wav2Vec2FeatureEncoder(nn.Module):392 """Construct the features from raw audio waveform"""393 394 def __init__(self, config):395 super().__init__()396 397 if config.feat_extract_norm == "group":398 conv_layers = [Wav2Vec2GroupNormConvLayer(config, layer_id=0)] + [399 Wav2Vec2NoLayerNormConvLayer(config, layer_id=i + 1) for i in range(config.num_feat_extract_layers - 1)400 ]401 elif config.feat_extract_norm == "layer":402 conv_layers = [403 Wav2Vec2LayerNormConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers)404 ]405 else:406 raise ValueError(407 f"`config.feat_extract_norm` is {config.feat_extract_norm}, but has to be one of ['group', 'layer']"408 )409 self.conv_layers = nn.ModuleList(conv_layers)410 self.gradient_checkpointing = False411 self._requires_grad = True412 413 def _freeze_parameters(self):414 for param in self.parameters():415 param.requires_grad = False416 self._requires_grad = False417 418 def forward(self, input_values):419 hidden_states = input_values[:, None]420 421 # make sure hidden_states require grad for gradient_checkpointing422 if self._requires_grad and self.training:423 hidden_states.requires_grad = True424 425 for conv_layer in self.conv_layers:426 hidden_states = conv_layer(hidden_states)427 428 return hidden_states429 430 431class Wav2Vec2FeatureExtractor(Wav2Vec2FeatureEncoder):432 def __init__(self, config):433 super().__init__(config)434 warnings.warn(435 f"The class `{self.__class__.__name__}` has been depreciated "436 "and will be removed in Transformers v5. "437 f"Use `{self.__class__.__bases__[0].__name__}` instead.",438 FutureWarning,439 )440 441 442class Wav2Vec2FeatureProjection(nn.Module):443 def __init__(self, config):444 super().__init__()445 self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps)446 self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size)447 self.dropout = nn.Dropout(config.feat_proj_dropout)448 449 def forward(self, hidden_states):450 # non-projected hidden states are needed for quantization451 norm_hidden_states = self.layer_norm(hidden_states)452 hidden_states = self.projection(norm_hidden_states)453 hidden_states = self.dropout(hidden_states)454 return hidden_states, norm_hidden_states455 456 457# Copied from transformers.models.bart.modeling_bart.eager_attention_forward458def eager_attention_forward(459 module: nn.Module,460 query: torch.Tensor,461 key: torch.Tensor,462 value: torch.Tensor,463 attention_mask: Optional[torch.Tensor],464 scaling: Optional[float] = None,465 dropout: float = 0.0,466 head_mask: Optional[torch.Tensor] = None,467 **kwargs,468):469 if scaling is None:470 scaling = query.size(-1) ** -0.5471 472 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling473 if attention_mask is not None:474 attn_weights = attn_weights + attention_mask475 476 attn_weights = nn.functional.softmax(attn_weights, dim=-1)477 478 if head_mask is not None:479 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)480 481 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)482 attn_output = torch.matmul(attn_weights, value)483 attn_output = attn_output.transpose(1, 2).contiguous()484 485 return attn_output, attn_weights486 487 488class Wav2Vec2Attention(nn.Module):489 """Multi-headed attention from 'Attention Is All You Need' paper"""490 491 def __init__(492 self,493 embed_dim: int,494 num_heads: int,495 dropout: float = 0.0,496 is_decoder: bool = False,497 bias: bool = True,498 is_causal: bool = False,499 config: Optional[Wav2Vec2Config] = None,500 ):501 super().__init__()502 self.embed_dim = embed_dim503 self.num_heads = num_heads504 self.dropout = dropout505 self.head_dim = embed_dim // num_heads506 self.config = config507 508 if (self.head_dim * num_heads) != self.embed_dim:509 raise ValueError(510 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"511 f" and `num_heads`: {num_heads})."512 )513 self.scaling = self.head_dim**-0.5514 self.is_decoder = is_decoder515 self.is_causal = is_causal516 517 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)518 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)519 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)520 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)521 522 def forward(523 self,524 hidden_states: torch.Tensor,525 key_value_states: Optional[torch.Tensor] = None,526 attention_mask: Optional[torch.Tensor] = None,527 layer_head_mask: Optional[torch.Tensor] = None,528 output_attentions: Optional[bool] = False,529 # TODO: we need a refactor so that the different attention modules can get their specific kwargs530 # ATM, we have mixed things encoder, decoder, and encoder-decoder attn531 **kwargs: Unpack[FlashAttentionKwargs],532 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:533 """Input shape: Batch x Time x Channel"""534 535 # if key_value_states are provided this layer is used as a cross-attention layer536 # for the decoder537 is_cross_attention = key_value_states is not None538 539 # determine input shapes540 bsz, tgt_len = hidden_states.shape[:-1]541 src_len = key_value_states.shape[1] if is_cross_attention else tgt_len542 543 q_input_shape = (bsz, tgt_len, -1, self.head_dim)544 kv_input_shape = (bsz, src_len, -1, self.head_dim)545 546 # get query proj547 query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)548 549 current_states = key_value_states if is_cross_attention else hidden_states550 key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2)551 value_states = self.v_proj(current_states).view(*kv_input_shape).transpose(1, 2)552 553 attention_interface: Callable = eager_attention_forward554 if self.config._attn_implementation != "eager":555 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]556 557 attn_output, attn_weights = attention_interface(558 self,559 query_states,560 key_states,561 value_states,562 attention_mask,563 dropout=0.0 if not self.training else self.dropout,564 scaling=self.scaling,565 output_attentions=output_attentions,566 head_mask=layer_head_mask,567 **kwargs,568 )569 570 attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()571 attn_output = self.out_proj(attn_output)572 573 return attn_output, attn_weights, None574 575 576class Wav2Vec2FeedForward(nn.Module):577 def __init__(self, config):578 super().__init__()579 self.intermediate_dropout = nn.Dropout(config.activation_dropout)580 581 self.intermediate_dense = nn.Linear(config.hidden_size, config.intermediate_size)582 if isinstance(config.hidden_act, str):583 self.intermediate_act_fn = ACT2FN[config.hidden_act]584 else:585 self.intermediate_act_fn = config.hidden_act586 587 self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size)588 self.output_dropout = nn.Dropout(config.hidden_dropout)589 590 def forward(self, hidden_states):591 hidden_states = self.intermediate_dense(hidden_states)592 hidden_states = self.intermediate_act_fn(hidden_states)593 hidden_states = self.intermediate_dropout(hidden_states)594 595 hidden_states = self.output_dense(hidden_states)596 hidden_states = self.output_dropout(hidden_states)597 return hidden_states598 599 600class Wav2Vec2EncoderLayer(GradientCheckpointingLayer):601 def __init__(self, config):602 super().__init__()603 self.attention = Wav2Vec2Attention(604 embed_dim=config.hidden_size,605 num_heads=config.num_attention_heads,606 dropout=config.attention_dropout,607 is_decoder=False,608 config=config,609 )610 611 self.dropout = nn.Dropout(config.hidden_dropout)612 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)613 self.feed_forward = Wav2Vec2FeedForward(config)614 self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)615 616 def forward(self, hidden_states, attention_mask=None, output_attentions=False):617 attn_residual = hidden_states618 hidden_states, attn_weights, _ = self.attention(619 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions620 )621 hidden_states = self.dropout(hidden_states)622 hidden_states = attn_residual + hidden_states623 624 hidden_states = self.layer_norm(hidden_states)625 hidden_states = hidden_states + self.feed_forward(hidden_states)626 hidden_states = self.final_layer_norm(hidden_states)627 628 outputs = (hidden_states,)629 630 if output_attentions:631 outputs += (attn_weights,)632 633 return outputs634 635 636class Wav2Vec2EncoderLayerStableLayerNorm(GradientCheckpointingLayer):637 def __init__(self, config):638 super().__init__()639 self.attention = Wav2Vec2Attention(640 embed_dim=config.hidden_size,641 num_heads=config.num_attention_heads,642 dropout=config.attention_dropout,643 is_decoder=False,644 config=config,645 )646 self.dropout = nn.Dropout(config.hidden_dropout)647 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)648 self.feed_forward = Wav2Vec2FeedForward(config)649 self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)650 651 if getattr(config, "adapter_attn_dim", None) is not None:652 self.adapter_layer = Wav2Vec2AttnAdapterLayer(config)653 else:654 self.adapter_layer = None655 656 def forward(657 self,658 hidden_states: torch.Tensor,659 attention_mask: Optional[torch.Tensor] = None,660 output_attentions: bool = False,661 ):662 attn_residual = hidden_states663 hidden_states = self.layer_norm(hidden_states)664 hidden_states, attn_weights, _ = self.attention(665 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions666 )667 hidden_states = self.dropout(hidden_states)668 hidden_states = attn_residual + hidden_states669 hidden_states = hidden_states + self.feed_forward(self.final_layer_norm(hidden_states))670 671 if self.adapter_layer is not None:672 hidden_states = hidden_states + self.adapter_layer(hidden_states)673 674 outputs = (hidden_states,)675 676 if output_attentions:677 outputs += (attn_weights,)678 679 return outputs680 681 682class Wav2Vec2Encoder(nn.Module):683 def __init__(self, config):684 super().__init__()685 self.config = config686 self.pos_conv_embed = Wav2Vec2PositionalConvEmbedding(config)687 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)688 self.dropout = nn.Dropout(config.hidden_dropout)689 self.layers = nn.ModuleList([Wav2Vec2EncoderLayer(config) for _ in range(config.num_hidden_layers)])690 self.gradient_checkpointing = False691 692 def forward(693 self,694 hidden_states: torch.tensor,695 attention_mask: Optional[torch.Tensor] = None,696 output_attentions: bool = False,697 output_hidden_states: bool = False,698 return_dict: bool = True,699 ):700 all_hidden_states = () if output_hidden_states else None701 all_self_attentions = () if output_attentions else None702 703 if attention_mask is not None:704 # make sure padded tokens output 0705 expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])706 hidden_states[~expand_attention_mask] = 0707 708 attention_mask = self._update_full_mask(709 attention_mask,710 hidden_states,711 )712 713 position_embeddings = self.pos_conv_embed(hidden_states)714 hidden_states = hidden_states + position_embeddings715 hidden_states = self.layer_norm(hidden_states)716 hidden_states = self.dropout(hidden_states)717 718 synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)719 720 for layer in self.layers:721 if output_hidden_states:722 all_hidden_states = all_hidden_states + (hidden_states,)723 724 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)725 dropout_probability = torch.rand([])726 727 skip_the_layer = self.training and dropout_probability < self.config.layerdrop728 if not skip_the_layer or synced_gpus:729 # under fsdp or deepspeed zero3 all gpus must run in sync730 layer_outputs = layer(731 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions732 )733 hidden_states = layer_outputs[0]734 735 if skip_the_layer:736 layer_outputs = (None, None)737 738 if output_attentions:739 all_self_attentions = all_self_attentions + (layer_outputs[1],)740 741 if output_hidden_states:742 all_hidden_states = all_hidden_states + (hidden_states,)743 744 if not return_dict:745 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)746 return BaseModelOutput(747 last_hidden_state=hidden_states,748 hidden_states=all_hidden_states,749 attentions=all_self_attentions,750 )751 752 # Copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_full_mask753 def _update_full_mask(754 self,755 attention_mask: Union[torch.Tensor, None],756 inputs_embeds: torch.Tensor,757 ):758 if attention_mask is not None:759 if self.config._attn_implementation == "flash_attention_2":760 attention_mask = attention_mask if 0 in attention_mask else None761 elif self.config._attn_implementation == "sdpa":762 # output_attentions=True & head_mask can not be supported when using SDPA, fall back to763 # the manual implementation that requires a 4D causal mask in all cases.764 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]765 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)766 elif self.config._attn_implementation == "flex_attention":767 if isinstance(attention_mask, torch.Tensor):768 attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False)769 else:770 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]771 attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)772 773 return attention_mask774 775 776class Wav2Vec2EncoderStableLayerNorm(nn.Module):777 def __init__(self, config):778 super().__init__()779 self.config = config780 self.pos_conv_embed = Wav2Vec2PositionalConvEmbedding(config)781 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)782 self.dropout = nn.Dropout(config.hidden_dropout)783 self.layers = nn.ModuleList(784 [Wav2Vec2EncoderLayerStableLayerNorm(config) for _ in range(config.num_hidden_layers)]785 )786 self.gradient_checkpointing = False787 788 def forward(789 self,790 hidden_states,791 attention_mask=None,792 output_attentions=False,793 output_hidden_states=False,794 return_dict=True,795 ):796 all_hidden_states = () if output_hidden_states else None797 all_self_attentions = () if output_attentions else None798 799 if attention_mask is not None:800 # make sure padded tokens output 0801 expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])802 hidden_states[~expand_attention_mask] = 0803 804 attention_mask = self._update_full_mask(805 attention_mask,806 hidden_states,807 )808 809 position_embeddings = self.pos_conv_embed(hidden_states)810 hidden_states = hidden_states + position_embeddings811 hidden_states = self.dropout(hidden_states)812 813 synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)814 815 for layer in self.layers:816 if output_hidden_states:817 all_hidden_states = all_hidden_states + (hidden_states,)818 819 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)820 dropout_probability = torch.rand([])821 822 skip_the_layer = self.training and dropout_probability < self.config.layerdrop823 if not skip_the_layer or synced_gpus:824 # under fsdp or deepspeed zero3 all gpus must run in sync825 # XXX: could optimize this like synced_gpus in generate_utils but not sure if it's worth the code complication826 layer_outputs = layer(827 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions828 )829 hidden_states = layer_outputs[0]830 831 if skip_the_layer:832 layer_outputs = (None, None)833 834 if output_attentions:835 all_self_attentions = all_self_attentions + (layer_outputs[1],)836 837 hidden_states = self.layer_norm(hidden_states)838 839 if output_hidden_states:840 all_hidden_states = all_hidden_states + (hidden_states,)841 842 if not return_dict:843 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)844 return BaseModelOutput(845 last_hidden_state=hidden_states,846 hidden_states=all_hidden_states,847 attentions=all_self_attentions,848 )849 850 # Copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_full_mask851 def _update_full_mask(852 self,853 attention_mask: Union[torch.Tensor, None],854 inputs_embeds: torch.Tensor,855 ):856 if attention_mask is not None:857 if self.config._attn_implementation == "flash_attention_2":858 attention_mask = attention_mask if 0 in attention_mask else None859 elif self.config._attn_implementation == "sdpa":860 # output_attentions=True & head_mask can not be supported when using SDPA, fall back to861 # the manual implementation that requires a 4D causal mask in all cases.862 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]863 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)864 elif self.config._attn_implementation == "flex_attention":865 if isinstance(attention_mask, torch.Tensor):866 attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False)867 else:868 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]869 attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)870 871 return attention_mask872 873 874class Wav2Vec2GumbelVectorQuantizer(nn.Module):875 """876 Vector quantization using gumbel softmax. See `[CATEGORICAL REPARAMETERIZATION WITH877 GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) for more information.878 """879 880 def __init__(self, config):881 super().__init__()882 self.num_groups = config.num_codevector_groups883 self.num_vars = config.num_codevectors_per_group884 885 if config.codevector_dim % self.num_groups != 0:886 raise ValueError(887 f"`config.codevector_dim {config.codevector_dim} must be divisible "888 f"by `config.num_codevector_groups` {self.num_groups} for concatenation"889 )890 891 # storage for codebook variables (codewords)892 self.codevectors = nn.Parameter(893 torch.FloatTensor(1, self.num_groups * self.num_vars, config.codevector_dim // self.num_groups)894 )895 self.weight_proj = nn.Linear(config.conv_dim[-1], self.num_groups * self.num_vars)896 897 # can be decayed for training898 self.temperature = 2899 900 @staticmethod901 def _compute_perplexity(probs, mask=None):902 if mask is not None:903 mask_extended = mask.flatten()[:, None, None].expand(probs.shape)904 probs = torch.where(mask_extended, probs, torch.zeros_like(probs))905 marginal_probs = probs.sum(dim=0) / mask.sum()906 else:907 marginal_probs = probs.mean(dim=0)908 909 perplexity = torch.exp(-torch.sum(marginal_probs * torch.log(marginal_probs + 1e-7), dim=-1)).sum()910 return perplexity911 912 def forward(self, hidden_states, mask_time_indices=None):913 batch_size, sequence_length, hidden_size = hidden_states.shape914 915 # project to codevector dim916 hidden_states = self.weight_proj(hidden_states)917 hidden_states = hidden_states.view(batch_size * sequence_length * self.num_groups, -1)918 919 if self.training:920 # sample code vector probs via gumbel in differentiateable way921 codevector_probs = nn.functional.gumbel_softmax(922 hidden_states.float(), tau=self.temperature, hard=True923 ).type_as(hidden_states)924 925 # compute perplexity926 codevector_soft_dist = torch.softmax(927 hidden_states.view(batch_size * sequence_length, self.num_groups, -1).float(), dim=-1928 )929 perplexity = self._compute_perplexity(codevector_soft_dist, mask_time_indices)930 else:931 # take argmax in non-differentiable way932 # comptute hard codevector distribution (one hot)933 codevector_idx = hidden_states.argmax(dim=-1)934 codevector_probs = hidden_states.new_zeros(hidden_states.shape).scatter_(935 -1, codevector_idx.view(-1, 1), 1.0936 )937 codevector_probs = codevector_probs.view(batch_size * sequence_length, self.num_groups, -1)938 939 perplexity = self._compute_perplexity(codevector_probs, mask_time_indices)940 941 codevector_probs = codevector_probs.view(batch_size * sequence_length, -1)942 # use probs to retrieve codevectors943 codevectors_per_group = codevector_probs.unsqueeze(-1) * self.codevectors944 codevectors = codevectors_per_group.view(batch_size * sequence_length, self.num_groups, self.num_vars, -1)945 codevectors = codevectors.sum(-2).view(batch_size, sequence_length, -1)946 947 return codevectors, perplexity948 949 950class Wav2Vec2Adapter(nn.Module):951 def __init__(self, config):952 super().__init__()953 954 # feature dim might need to be down-projected955 if config.output_hidden_size != config.hidden_size:956 self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)957 self.proj_layer_norm = nn.LayerNorm(config.output_hidden_size)958 else:959 self.proj = self.proj_layer_norm = None960 961 self.layers = nn.ModuleList(Wav2Vec2AdapterLayer(config) for _ in range(config.num_adapter_layers))962 self.layerdrop = config.layerdrop963 964 def forward(self, hidden_states):965 # down project hidden_states if necessary966 if self.proj is not None and self.proj_layer_norm is not None:967 hidden_states = self.proj(hidden_states)968 hidden_states = self.proj_layer_norm(hidden_states)969 970 hidden_states = hidden_states.transpose(1, 2)971 972 for layer in self.layers:973 layerdrop_prob = np.random.random()974 if not self.training or (layerdrop_prob > self.layerdrop):975 hidden_states = layer(hidden_states)976 977 hidden_states = hidden_states.transpose(1, 2)978 return hidden_states979 980 981class Wav2Vec2AdapterLayer(nn.Module):982 def __init__(self, config):983 super().__init__()984 self.conv = nn.Conv1d(985 config.output_hidden_size,986 2 * config.output_hidden_size,987 config.adapter_kernel_size,988 stride=config.adapter_stride,989 padding=1,990 )991 992 def forward(self, hidden_states):993 hidden_states = self.conv(hidden_states)994 hidden_states = nn.functional.glu(hidden_states, dim=1)995 996 return hidden_states997 998 999class Wav2Vec2AttnAdapterLayer(nn.Module):1000 def __init__(self, config):1001 """1002 Implements adapter modules directly with 3D tensor weight as parameters and without using ModuleList to speed1003 up training throughput.1004 """1005 super().__init__()1006 self.input_dim = config.adapter_attn_dim1007 self.hidden_dim = config.hidden_size1008 1009 self.norm = nn.LayerNorm(self.hidden_dim)1010 self.linear_1 = nn.Linear(self.hidden_dim, self.input_dim)1011 self.act_fn = nn.ReLU()1012 self.linear_2 = nn.Linear(self.input_dim, self.hidden_dim)1013 1014 def forward(self, hidden_states: torch.FloatTensor):1015 hidden_states = self.norm(hidden_states)1016 1017 hidden_states = self.linear_1(hidden_states)1018 hidden_states = self.act_fn(hidden_states)1019 hidden_states = self.linear_2(hidden_states)1020 1021 return hidden_states1022 1023 1024@auto_docstring1025class Wav2Vec2PreTrainedModel(PreTrainedModel):1026 config: Wav2Vec2Config1027 base_model_prefix = "wav2vec2"1028 main_input_name = "input_values"1029 supports_gradient_checkpointing = True1030 _supports_flash_attn = True1031 _supports_sdpa = True1032 _supports_flex_attn = True1033 1034 def _init_weights(self, module):1035 """Initialize the weights"""1036 # Wav2Vec2ForPreTraining last 2 linear layers need standard Linear init.1037 if isinstance(module, Wav2Vec2ForPreTraining):1038 module.project_hid.reset_parameters()1039 module.project_q.reset_parameters()1040 module.project_hid._is_hf_initialized = True1041 module.project_q._is_hf_initialized = True1042 # gumbel softmax requires special init1043 elif isinstance(module, Wav2Vec2GumbelVectorQuantizer):1044 module.weight_proj.weight.data.normal_(mean=0.0, std=1)1045 module.weight_proj.bias.data.zero_()1046 nn.init.uniform_(module.codevectors)1047 elif isinstance(module, Wav2Vec2PositionalConvEmbedding):1048 nn.init.normal_(1049 module.conv.weight,1050 mean=0,1051 std=2 * math.sqrt(1 / (module.conv.kernel_size[0] * module.conv.in_channels)),1052 )1053 nn.init.constant_(module.conv.bias, 0)1054 elif isinstance(module, Wav2Vec2FeatureProjection):1055 k = math.sqrt(1 / module.projection.in_features)1056 nn.init.uniform_(module.projection.weight, a=-k, b=k)1057 nn.init.uniform_(module.projection.bias, a=-k, b=k)1058 elif isinstance(module, nn.Linear):1059 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)1060 1061 if module.bias is not None:1062 module.bias.data.zero_()1063 elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):1064 module.bias.data.zero_()1065 module.weight.data.fill_(1.0)1066 elif isinstance(module, nn.Conv1d):1067 nn.init.kaiming_normal_(module.weight)1068 1069 if module.bias is not None:1070 k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))1071 nn.init.uniform_(module.bias, a=-k, b=k)1072 1073 def _get_feat_extract_output_lengths(1074 self, input_lengths: Union[torch.LongTensor, int], add_adapter: Optional[bool] = None1075 ):1076 """1077 Computes the output length of the convolutional layers1078 """1079 1080 add_adapter = self.config.add_adapter if add_adapter is None else add_adapter1081 1082 def _conv_out_length(input_length, kernel_size, stride):1083 # 1D convolutional layer output length formula taken1084 # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html1085 return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 11086 1087 for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):1088 input_lengths = _conv_out_length(input_lengths, kernel_size, stride)1089 1090 if add_adapter:1091 for _ in range(self.config.num_adapter_layers):1092 input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)1093 1094 return input_lengths1095 1096 def _get_feature_vector_attention_mask(1097 self, feature_vector_length: int, attention_mask: torch.LongTensor, add_adapter=None1098 ):1099 # Effectively attention_mask.sum(-1), but not inplace to be able to run1100 # on inference mode.1101 non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1]1102 1103 output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)1104 output_lengths = output_lengths.to(torch.long)1105 1106 batch_size = attention_mask.shape[0]1107 1108 attention_mask = torch.zeros(1109 (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device1110 )1111 # these two operations makes sure that all values before the output lengths idxs are attended to1112 attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 11113 attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()1114 return attention_mask1115 1116 def _get_adapters(self):1117 if self.config.adapter_attn_dim is None:1118 raise ValueError(f"{self.__class__} has no adapter layers. Make sure to define `config.adapter_attn_dim`.")1119 1120 adapter_weights = {}1121 for name, module in self.named_modules():1122 if isinstance(module, Wav2Vec2AttnAdapterLayer):1123 for param_name, param in module.named_parameters():1124 adapter_weights[".".join([name, param_name])] = param1125 1126 if isinstance(self, Wav2Vec2ForCTC):1127 for name, param in self.lm_head.named_parameters():1128 adapter_weights[".".join(["lm_head", name])] = param1129 1130 return adapter_weights1131 1132 def init_adapter_layers(self):1133 """1134 (Re-)initialize attention adapter layers and lm head for adapter-only fine-tuning1135 """1136 # init attention adapters1137 for module in self.modules():1138 if isinstance(module, Wav2Vec2AttnAdapterLayer):1139 self._init_weights(module)1140 1141 # init lm head1142 if isinstance(self, Wav2Vec2ForCTC):1143 self._init_weights(self.lm_head)1144 1145 def load_adapter(self, target_lang: str, force_load=True, **kwargs):1146 r"""1147 Load a language adapter model from a pre-trained adapter model.1148 1149 Parameters:1150 target_lang (`str`):1151 Has to be a language id of an existing adapter weight. Adapter weights are stored in the format1152 adapter.<lang>.safetensors or adapter.<lang>.bin1153 force_load (`bool`, defaults to `True`):1154 Whether the weights shall be loaded even if `target_lang` matches `self.target_lang`.1155 cache_dir (`Union[str, os.PathLike]`, *optional*):1156 Path to a directory in which a downloaded pretrained model configuration should be cached if the1157 standard cache should not be used.1158 force_download (`bool`, *optional*, defaults to `False`):1159 Whether or not to force the (re-)download of the model weights and configuration files, overriding the1160 cached versions if they exist.1161 resume_download:1162 Deprecated and ignored. All downloads are now resumed by default when possible.1163 Will be removed in v5 of Transformers.1164 proxies (`dict[str, str]`, *optional*):1165 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',1166 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.1167 local_files_only(`bool`, *optional*, defaults to `False`):1168 Whether or not to only look at local files (i.e., do not try to download the model).1169 token (`str` or `bool`, *optional*):1170 The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use1171 the token generated when running `hf auth login` (stored in `~/.huggingface`).1172 revision (`str`, *optional*, defaults to `"main"`):1173 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a1174 git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any1175 identifier allowed by git.1176 1177 <Tip>1178 1179 To test a pull request you made on the Hub, you can pass `revision="refs/pr/<pr_number>"`.1180 1181 </Tip>1182 1183 mirror (`str`, *optional*):1184 Mirror source to accelerate downloads in China. If you are from China and have an accessibility1185 problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.1186 Please refer to the mirror site for more information.1187 1188 <Tip>1189 1190 Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to1191 use this method in a firewalled environment.1192 1193 </Tip>1194 1195 Examples:1196 1197 ```python1198 >>> from transformers import Wav2Vec2ForCTC, AutoProcessor1199 1200 >>> ckpt = "facebook/mms-1b-all"