Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/data2vec/modular_data2vec_audio.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_data2vec_audio.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2022 The HuggingFace Inc. team.9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14# http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22import math23import warnings24from typing import Callable, Optional, Union25 26import numpy as np27import torch28from torch import nn29from torch.nn import CrossEntropyLoss30 31from ...activations import ACT2FN32from ...integrations.deepspeed import is_deepspeed_zero3_enabled33from ...integrations.fsdp import is_fsdp_managed_module34from ...modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa35from ...modeling_flash_attention_utils import FlashAttentionKwargs36from ...modeling_layers import GradientCheckpointingLayer37from ...modeling_outputs import (38 BaseModelOutput,39 CausalLMOutput,40 SequenceClassifierOutput,41 TokenClassifierOutput,42 Wav2Vec2BaseModelOutput,43 XVectorOutput,44)45from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel46from ...processing_utils import Unpack47from ...utils import auto_docstring, is_peft_available, is_torch_flex_attn_available48from .configuration_data2vec_audio import Data2VecAudioConfig49 50 51if is_torch_flex_attn_available():52 from ...integrations.flex_attention import make_flex_block_causal_mask53 54 55class Data2VecAudioConvLayer(GradientCheckpointingLayer):56 def __init__(self, config, layer_id=0):57 super().__init__()58 self.in_conv_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 159 self.out_conv_dim = config.conv_dim[layer_id]60 61 self.conv = nn.Conv1d(62 self.in_conv_dim,63 self.out_conv_dim,64 kernel_size=config.conv_kernel[layer_id],65 stride=config.conv_stride[layer_id],66 bias=config.conv_bias,67 )68 self.layer_norm = nn.LayerNorm(self.out_conv_dim, elementwise_affine=True)69 self.activation = ACT2FN[config.feat_extract_activation]70 71 def forward(self, hidden_states):72 hidden_states = self.conv(hidden_states)73 74 hidden_states = hidden_states.transpose(-2, -1)75 hidden_states = self.layer_norm(hidden_states)76 hidden_states = hidden_states.transpose(-2, -1)77 78 hidden_states = self.activation(hidden_states)79 return hidden_states80 81 82class Data2VecAudioPadLayer(nn.Module):83 def __init__(self, num_conv_pos_embeddings):84 super().__init__()85 self.num_pad_remove = 1 if num_conv_pos_embeddings % 2 == 0 else 086 87 def forward(self, hidden_states):88 if self.num_pad_remove > 0:89 hidden_states = hidden_states[:, :, : -self.num_pad_remove]90 return hidden_states91 92 93class Data2VecAudioPositionalConvLayer(nn.Module):94 def __init__(self, config):95 super().__init__()96 self.conv = nn.Conv1d(97 config.hidden_size,98 config.hidden_size,99 kernel_size=config.conv_pos_kernel_size,100 padding=config.conv_pos_kernel_size // 2,101 groups=config.num_conv_pos_embedding_groups,102 )103 104 self.padding = Data2VecAudioPadLayer(config.conv_pos_kernel_size)105 self.activation = ACT2FN[config.feat_extract_activation]106 # no learnable parameters107 self.layer_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False)108 109 def forward(self, hidden_states):110 hidden_states = self.conv(hidden_states)111 hidden_states = self.padding(hidden_states)112 113 hidden_states = hidden_states.transpose(1, 2)114 hidden_states = self.layer_norm(hidden_states)115 hidden_states = hidden_states.transpose(1, 2)116 hidden_states = self.activation(hidden_states)117 return hidden_states118 119 120class Data2VecAudioPositionalConvEmbedding(nn.Module):121 def __init__(self, config):122 super().__init__()123 self.layers = nn.ModuleList(124 [Data2VecAudioPositionalConvLayer(config) for _ in range(config.num_conv_pos_embeddings)]125 )126 127 def forward(self, hidden_states):128 hidden_states = hidden_states.transpose(1, 2)129 for layer in self.layers:130 hidden_states = layer(hidden_states)131 hidden_states = hidden_states.transpose(1, 2)132 return hidden_states133 134 135class Data2VecAudioFeatureEncoder(nn.Module):136 """Construct the features from raw audio waveform"""137 138 def __init__(self, config):139 super().__init__()140 self.conv_layers = nn.ModuleList(141 [Data2VecAudioConvLayer(config, layer_id=i) for i in range(config.num_feat_extract_layers)]142 )143 self.gradient_checkpointing = False144 self._requires_grad = True145 146 def _freeze_parameters(self):147 for param in self.parameters():148 param.requires_grad = False149 self._requires_grad = False150 151 def forward(self, input_values):152 hidden_states = input_values[:, None]153 154 # make sure hidden_states require grad for gradient_checkpointing155 if self._requires_grad and self.training:156 hidden_states.requires_grad = True157 158 for conv_layer in self.conv_layers:159 hidden_states = conv_layer(hidden_states)160 161 return hidden_states162 163 164class Data2VecAudioFeatureProjection(nn.Module):165 def __init__(self, config):166 super().__init__()167 self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps)168 self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size)169 self.dropout = nn.Dropout(config.feat_proj_dropout)170 171 def forward(self, hidden_states):172 # non-projected hidden states are needed for quantization173 norm_hidden_states = self.layer_norm(hidden_states)174 hidden_states = self.projection(norm_hidden_states)175 hidden_states = self.dropout(hidden_states)176 return hidden_states, norm_hidden_states177 178 179def eager_attention_forward(180 module: nn.Module,181 query: torch.Tensor,182 key: torch.Tensor,183 value: torch.Tensor,184 attention_mask: Optional[torch.Tensor],185 scaling: Optional[float] = None,186 dropout: float = 0.0,187 head_mask: Optional[torch.Tensor] = None,188 **kwargs,189):190 if scaling is None:191 scaling = query.size(-1) ** -0.5192 193 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling194 if attention_mask is not None:195 attn_weights = attn_weights + attention_mask196 197 attn_weights = nn.functional.softmax(attn_weights, dim=-1)198 199 if head_mask is not None:200 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)201 202 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)203 attn_output = torch.matmul(attn_weights, value)204 attn_output = attn_output.transpose(1, 2).contiguous()205 206 return attn_output, attn_weights207 208 209class Data2VecAudioAttention(nn.Module):210 """Multi-headed attention from 'Attention Is All You Need' paper"""211 212 def __init__(213 self,214 embed_dim: int,215 num_heads: int,216 dropout: float = 0.0,217 is_decoder: bool = False,218 bias: bool = True,219 is_causal: bool = False,220 config: Optional[Data2VecAudioConfig] = None,221 ):222 super().__init__()223 self.embed_dim = embed_dim224 self.num_heads = num_heads225 self.dropout = dropout226 self.head_dim = embed_dim // num_heads227 self.config = config228 229 if (self.head_dim * num_heads) != self.embed_dim:230 raise ValueError(231 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"232 f" and `num_heads`: {num_heads})."233 )234 self.scaling = self.head_dim**-0.5235 self.is_decoder = is_decoder236 self.is_causal = is_causal237 238 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)239 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)240 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)241 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)242 243 def forward(244 self,245 hidden_states: torch.Tensor,246 key_value_states: Optional[torch.Tensor] = None,247 attention_mask: Optional[torch.Tensor] = None,248 layer_head_mask: Optional[torch.Tensor] = None,249 output_attentions: Optional[bool] = False,250 # TODO: we need a refactor so that the different attention modules can get their specific kwargs251 # ATM, we have mixed things encoder, decoder, and encoder-decoder attn252 **kwargs: Unpack[FlashAttentionKwargs],253 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:254 """Input shape: Batch x Time x Channel"""255 256 # if key_value_states are provided this layer is used as a cross-attention layer257 # for the decoder258 is_cross_attention = key_value_states is not None259 260 # determine input shapes261 bsz, tgt_len = hidden_states.shape[:-1]262 src_len = key_value_states.shape[1] if is_cross_attention else tgt_len263 264 q_input_shape = (bsz, tgt_len, -1, self.head_dim)265 kv_input_shape = (bsz, src_len, -1, self.head_dim)266 267 # get query proj268 query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)269 270 current_states = key_value_states if is_cross_attention else hidden_states271 key_states = self.k_proj(current_states).view(*kv_input_shape).transpose(1, 2)272 value_states = self.v_proj(current_states).view(*kv_input_shape).transpose(1, 2)273 274 attention_interface: Callable = eager_attention_forward275 if self.config._attn_implementation != "eager":276 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]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.dropout,285 scaling=self.scaling,286 output_attentions=output_attentions,287 head_mask=layer_head_mask,288 **kwargs,289 )290 291 attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()292 attn_output = self.out_proj(attn_output)293 294 return attn_output, attn_weights, None295 296 297class Data2VecAudioFeedForward(nn.Module):298 def __init__(self, config):299 super().__init__()300 self.intermediate_dropout = nn.Dropout(config.activation_dropout)301 302 self.intermediate_dense = nn.Linear(config.hidden_size, config.intermediate_size)303 if isinstance(config.hidden_act, str):304 self.intermediate_act_fn = ACT2FN[config.hidden_act]305 else:306 self.intermediate_act_fn = config.hidden_act307 308 self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size)309 self.output_dropout = nn.Dropout(config.hidden_dropout)310 311 def forward(self, hidden_states):312 hidden_states = self.intermediate_dense(hidden_states)313 hidden_states = self.intermediate_act_fn(hidden_states)314 hidden_states = self.intermediate_dropout(hidden_states)315 316 hidden_states = self.output_dense(hidden_states)317 hidden_states = self.output_dropout(hidden_states)318 return hidden_states319 320 321class Data2VecAudioEncoderLayer(GradientCheckpointingLayer):322 def __init__(self, config):323 super().__init__()324 self.attention = Data2VecAudioAttention(325 embed_dim=config.hidden_size,326 num_heads=config.num_attention_heads,327 dropout=config.attention_dropout,328 is_decoder=False,329 config=config,330 )331 332 self.dropout = nn.Dropout(config.hidden_dropout)333 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)334 self.feed_forward = Data2VecAudioFeedForward(config)335 self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)336 337 def forward(self, hidden_states, attention_mask=None, output_attentions=False):338 attn_residual = hidden_states339 hidden_states, attn_weights, _ = self.attention(340 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions341 )342 hidden_states = self.dropout(hidden_states)343 hidden_states = attn_residual + hidden_states344 345 hidden_states = self.layer_norm(hidden_states)346 hidden_states = hidden_states + self.feed_forward(hidden_states)347 hidden_states = self.final_layer_norm(hidden_states)348 349 outputs = (hidden_states,)350 351 if output_attentions:352 outputs += (attn_weights,)353 354 return outputs355 356 357class Data2VecAudioEncoder(nn.Module):358 def __init__(self, config):359 super().__init__()360 self.config = config361 self.pos_conv_embed = Data2VecAudioPositionalConvEmbedding(config)362 self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)363 self.dropout = nn.Dropout(config.hidden_dropout)364 self.layers = nn.ModuleList([Data2VecAudioEncoderLayer(config) for _ in range(config.num_hidden_layers)])365 self.gradient_checkpointing = False366 367 def forward(368 self,369 hidden_states: torch.tensor,370 attention_mask: Optional[torch.Tensor] = None,371 output_attentions: bool = False,372 output_hidden_states: bool = False,373 return_dict: bool = True,374 ):375 all_hidden_states = () if output_hidden_states else None376 all_self_attentions = () if output_attentions else None377 378 if attention_mask is not None:379 # make sure padded tokens output 0380 expand_attention_mask = attention_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])381 hidden_states[~expand_attention_mask] = 0382 383 attention_mask = self._update_full_mask(384 attention_mask,385 hidden_states,386 )387 388 position_embeddings = self.pos_conv_embed(hidden_states)389 hidden_states = hidden_states + position_embeddings390 hidden_states = self.layer_norm(hidden_states)391 hidden_states = self.dropout(hidden_states)392 393 synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)394 395 for layer in self.layers:396 if output_hidden_states:397 all_hidden_states = all_hidden_states + (hidden_states,)398 399 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)400 dropout_probability = torch.rand([])401 402 skip_the_layer = self.training and dropout_probability < self.config.layerdrop403 if not skip_the_layer or synced_gpus:404 # under fsdp or deepspeed zero3 all gpus must run in sync405 layer_outputs = layer(406 hidden_states, attention_mask=attention_mask, output_attentions=output_attentions407 )408 hidden_states = layer_outputs[0]409 410 if skip_the_layer:411 layer_outputs = (None, None)412 413 if output_attentions:414 all_self_attentions = all_self_attentions + (layer_outputs[1],)415 416 if output_hidden_states:417 all_hidden_states = all_hidden_states + (hidden_states,)418 419 if not return_dict:420 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)421 return BaseModelOutput(422 last_hidden_state=hidden_states,423 hidden_states=all_hidden_states,424 attentions=all_self_attentions,425 )426 427 def _update_full_mask(428 self,429 attention_mask: Union[torch.Tensor, None],430 inputs_embeds: torch.Tensor,431 ):432 if attention_mask is not None:433 if self.config._attn_implementation == "flash_attention_2":434 attention_mask = attention_mask if 0 in attention_mask else None435 elif self.config._attn_implementation == "sdpa":436 # output_attentions=True & head_mask can not be supported when using SDPA, fall back to437 # the manual implementation that requires a 4D causal mask in all cases.438 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]439 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)440 elif self.config._attn_implementation == "flex_attention":441 if isinstance(attention_mask, torch.Tensor):442 attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False)443 else:444 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]445 attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)446 447 return attention_mask448 449 450class Data2VecAudioAdapterLayer(nn.Module):451 def __init__(self, config):452 super().__init__()453 self.conv = nn.Conv1d(454 config.output_hidden_size,455 2 * config.output_hidden_size,456 config.adapter_kernel_size,457 stride=config.adapter_stride,458 padding=1,459 )460 461 def forward(self, hidden_states):462 hidden_states = self.conv(hidden_states)463 hidden_states = nn.functional.glu(hidden_states, dim=1)464 465 return hidden_states466 467 468class Data2VecAudioAdapter(nn.Module):469 def __init__(self, config):470 super().__init__()471 472 # feature dim might need to be down-projected473 if config.output_hidden_size != config.hidden_size:474 self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)475 self.proj_layer_norm = nn.LayerNorm(config.output_hidden_size)476 else:477 self.proj = self.proj_layer_norm = None478 479 self.layers = nn.ModuleList(Data2VecAudioAdapterLayer(config) for _ in range(config.num_adapter_layers))480 self.layerdrop = config.layerdrop481 482 def forward(self, hidden_states):483 # down project hidden_states if necessary484 if self.proj is not None and self.proj_layer_norm is not None:485 hidden_states = self.proj(hidden_states)486 hidden_states = self.proj_layer_norm(hidden_states)487 488 hidden_states = hidden_states.transpose(1, 2)489 490 for layer in self.layers:491 layerdrop_prob = np.random.random()492 if not self.training or (layerdrop_prob > self.layerdrop):493 hidden_states = layer(hidden_states)494 495 hidden_states = hidden_states.transpose(1, 2)496 return hidden_states497 498 499@auto_docstring500class Data2VecAudioPreTrainedModel(PreTrainedModel):501 config: Data2VecAudioConfig502 base_model_prefix = "data2vec_audio"503 main_input_name = "input_values"504 supports_gradient_checkpointing = True505 _supports_flash_attn = True506 _supports_sdpa = True507 _supports_flex_attn = True508 509 def _init_weights(self, module):510 """Initialize the weights"""511 if isinstance(module, Data2VecAudioFeatureProjection):512 k = math.sqrt(1 / module.projection.in_features)513 nn.init.uniform_(module.projection.weight, a=-k, b=k)514 nn.init.uniform_(module.projection.bias, a=-k, b=k)515 elif isinstance(module, Data2VecAudioPositionalConvLayer):516 nn.init.constant_(module.conv.bias, 0)517 elif isinstance(module, nn.Linear):518 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)519 520 if module.bias is not None:521 module.bias.data.zero_()522 elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):523 if module.bias is not None:524 module.bias.data.zero_()525 if module.weight is not None:526 module.weight.data.fill_(1.0)527 elif isinstance(module, nn.Conv1d):528 nn.init.kaiming_normal_(module.weight)529 530 if module.bias is not None:531 k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))532 nn.init.uniform_(module.bias, a=-k, b=k)533 534 def _get_feat_extract_output_lengths(535 self, input_lengths: Union[torch.LongTensor, int], add_adapter: Optional[bool] = None536 ):537 """538 Computes the output length of the convolutional layers539 """540 541 add_adapter = self.config.add_adapter if add_adapter is None else add_adapter542 543 def _conv_out_length(input_length, kernel_size, stride):544 # 1D convolutional layer output length formula taken545 # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html546 return torch.div(input_length - kernel_size, stride, rounding_mode="floor") + 1547 548 for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):549 input_lengths = _conv_out_length(input_lengths, kernel_size, stride)550 551 if add_adapter:552 for _ in range(self.config.num_adapter_layers):553 input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)554 555 return input_lengths556 557 def _get_feature_vector_attention_mask(558 self, feature_vector_length: int, attention_mask: torch.LongTensor, add_adapter=None559 ):560 # Effectively attention_mask.sum(-1), but not inplace to be able to run561 # on inference mode.562 non_padded_lengths = attention_mask.cumsum(dim=-1)[:, -1]563 564 output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)565 output_lengths = output_lengths.to(torch.long)566 567 batch_size = attention_mask.shape[0]568 569 attention_mask = torch.zeros(570 (batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device571 )572 # these two operations makes sure that all values before the output lengths idxs are attended to573 attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1574 attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()575 return attention_mask576 577 578def _compute_mask_indices(579 shape: tuple[int, int],580 mask_prob: float,581 mask_length: int,582 attention_mask: Optional[torch.LongTensor] = None,583 min_masks: int = 0,584) -> np.ndarray:585 """586 Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for587 ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on588 CPU as part of the preprocessing during training.589 590 Args:591 shape: The shape for which to compute masks. This should be of a tuple of size 2 where592 the first element is the batch size and the second element is the length of the axis to span.593 mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of594 independently generated mask spans of length `mask_length` is computed by595 `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the596 actual percentage will be smaller.597 mask_length: size of the mask598 min_masks: minimum number of masked spans599 attention_mask: A (right-padded) attention mask which independently shortens the feature axis of600 each batch dimension.601 """602 batch_size, sequence_length = shape603 604 if mask_length < 1:605 raise ValueError("`mask_length` has to be bigger than 0.")606 607 if mask_length > sequence_length:608 raise ValueError(609 f"`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: {mask_length}"610 f" and `sequence_length`: {sequence_length}`"611 )612 613 # epsilon is used for probabilistic rounding614 epsilon = np.random.rand(1).item()615 616 def compute_num_masked_span(input_length):617 """Given input length, compute how many spans should be masked"""618 num_masked_span = int(mask_prob * input_length / mask_length + epsilon)619 num_masked_span = max(num_masked_span, min_masks)620 621 # make sure num masked span <= sequence_length622 if num_masked_span * mask_length > sequence_length:623 num_masked_span = sequence_length // mask_length624 625 # make sure num_masked span is also <= input_length - (mask_length - 1)626 if input_length - (mask_length - 1) < num_masked_span:627 num_masked_span = max(input_length - (mask_length - 1), 0)628 629 return num_masked_span630 631 # compute number of masked spans in batch632 input_lengths = (633 attention_mask.detach().sum(-1).tolist()634 if attention_mask is not None635 else [sequence_length for _ in range(batch_size)]636 )637 638 # SpecAugment mask to fill639 spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)640 spec_aug_mask_idxs = []641 642 max_num_masked_span = compute_num_masked_span(sequence_length)643 644 if max_num_masked_span == 0:645 return spec_aug_mask646 647 for input_length in input_lengths:648 # compute num of masked spans for this input649 num_masked_span = compute_num_masked_span(input_length)650 651 # get random indices to mask652 spec_aug_mask_idx = np.random.choice(653 np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False654 )655 656 # pick first sampled index that will serve as a dummy index to pad vector657 # to ensure same dimension for all batches due to probabilistic rounding658 # Picking first sample just pads those vectors twice.659 if len(spec_aug_mask_idx) == 0:660 # this case can only happen if `input_length` is strictly smaller then661 # `sequence_length` in which case the last token has to be a padding662 # token which we can use as a dummy mask id663 dummy_mask_idx = sequence_length - 1664 else:665 dummy_mask_idx = spec_aug_mask_idx[0]666 667 spec_aug_mask_idx = np.concatenate(668 [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]669 )670 spec_aug_mask_idxs.append(spec_aug_mask_idx)671 672 spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)673 674 # expand masked indices to masked spans675 spec_aug_mask_idxs = np.broadcast_to(676 spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)677 )678 spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)679 680 # add offset to the starting indexes so that indexes now create a span681 offsets = np.arange(mask_length)[None, None, :]682 offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(683 batch_size, max_num_masked_span * mask_length684 )685 spec_aug_mask_idxs = spec_aug_mask_idxs + offsets686 687 # ensure that we cannot have indices larger than sequence_length688 if spec_aug_mask_idxs.max() > sequence_length - 1:689 spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1690 691 # scatter indices to mask692 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)693 694 return spec_aug_mask695 696 697Data2VecAudioBaseModelOutput = Wav2Vec2BaseModelOutput698 699 700@auto_docstring701class Data2VecAudioModel(Data2VecAudioPreTrainedModel):702 def __init__(self, config: Data2VecAudioConfig):703 super().__init__(config)704 self.config = config705 self.feature_extractor = Data2VecAudioFeatureEncoder(config)706 self.feature_projection = Data2VecAudioFeatureProjection(config)707 708 # model only needs masking vector if mask prob is > 0.0709 if config.mask_time_prob > 0.0 or config.mask_feature_prob > 0.0:710 self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_())711 712 self.encoder = Data2VecAudioEncoder(config)713 714 self.adapter = Data2VecAudioAdapter(config) if config.add_adapter else None715 716 # Initialize weights and apply final processing717 self.post_init()718 719 def freeze_feature_encoder(self):720 """721 Calling this function will disable the gradient computation for the feature encoder so that its parameter will722 not be updated during training.723 """724 self.feature_extractor._freeze_parameters()725 726 def _mask_hidden_states(727 self,728 hidden_states: torch.FloatTensor,729 mask_time_indices: Optional[torch.FloatTensor] = None,730 attention_mask: Optional[torch.LongTensor] = None,731 ):732 """733 Masks extracted features along time axis and/or along feature axis according to734 [SpecAugment](https://huggingface.co/papers/1904.08779).735 """736 737 # `config.apply_spec_augment` can set masking to False738 if not getattr(self.config, "apply_spec_augment", True):739 return hidden_states740 741 # generate indices & apply SpecAugment along time axis742 batch_size, sequence_length, hidden_size = hidden_states.size()743 744 if mask_time_indices is not None:745 # apply SpecAugment along time axis with given mask_time_indices746 hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)747 elif self.config.mask_time_prob > 0 and self.training:748 mask_time_indices = _compute_mask_indices(749 (batch_size, sequence_length),750 mask_prob=self.config.mask_time_prob,751 mask_length=self.config.mask_time_length,752 attention_mask=attention_mask,753 min_masks=self.config.mask_time_min_masks,754 )755 mask_time_indices = torch.tensor(mask_time_indices, device=hidden_states.device, dtype=torch.bool)756 hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)757 758 if self.config.mask_feature_prob > 0 and self.training:759 # generate indices & apply SpecAugment along feature axis760 mask_feature_indices = _compute_mask_indices(761 (batch_size, hidden_size),762 mask_prob=self.config.mask_feature_prob,763 mask_length=self.config.mask_feature_length,764 min_masks=self.config.mask_feature_min_masks,765 )766 mask_feature_indices = torch.tensor(mask_feature_indices, device=hidden_states.device, dtype=torch.bool)767 mask_feature_indices = mask_feature_indices[:, None].expand(-1, sequence_length, -1)768 hidden_states[mask_feature_indices] = 0769 770 return hidden_states771 772 @auto_docstring773 def forward(774 self,775 input_values: Optional[torch.Tensor],776 attention_mask: Optional[torch.Tensor] = None,777 mask_time_indices: Optional[torch.FloatTensor] = None,778 output_attentions: Optional[bool] = None,779 output_hidden_states: Optional[bool] = None,780 return_dict: Optional[bool] = None,781 ) -> Union[tuple, Data2VecAudioBaseModelOutput]:782 r"""783 mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):784 Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict785 masked extracted features in *config.proj_codevector_dim* space.786 """787 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions788 output_hidden_states = (789 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states790 )791 return_dict = return_dict if return_dict is not None else self.config.use_return_dict792 793 extract_features = self.feature_extractor(input_values)794 extract_features = extract_features.transpose(1, 2)795 796 if attention_mask is not None:797 # compute reduced attention_mask corresponding to feature vectors798 attention_mask = self._get_feature_vector_attention_mask(799 extract_features.shape[1], attention_mask, add_adapter=False800 )801 802 hidden_states, extract_features = self.feature_projection(extract_features)803 hidden_states = self._mask_hidden_states(804 hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask805 )806 807 encoder_outputs = self.encoder(808 hidden_states,809 attention_mask=attention_mask,810 output_attentions=output_attentions,811 output_hidden_states=output_hidden_states,812 return_dict=return_dict,813 )814 815 hidden_states = encoder_outputs[0]816 817 if self.adapter is not None:818 hidden_states = self.adapter(hidden_states)819 820 if not return_dict:821 return (hidden_states, extract_features) + encoder_outputs[1:]822 823 return Data2VecAudioBaseModelOutput(824 last_hidden_state=hidden_states,825 extract_features=extract_features,826 hidden_states=encoder_outputs.hidden_states,827 attentions=encoder_outputs.attentions,828 )829 830 831_HIDDEN_STATES_START_POSITION = 2832 833 834@auto_docstring(835 custom_intro="""836 Data2VecAudio Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).837 """838)839class Data2VecAudioForCTC(Data2VecAudioPreTrainedModel):840 def __init__(self, config):841 r"""842 target_lang (`str`, *optional*):843 Language id of adapter weights. Adapter weights are stored in the format adapter.<lang>.safetensors or844 adapter.<lang>.bin. Only relevant when using an instance of [`Data2VecAudioForCTC`] with adapters. Uses 'eng' by845 default.846 """847 super().__init__(config)848 849 self.data2vec_audio = Data2VecAudioModel(config)850 self.dropout = nn.Dropout(config.final_dropout)851 852 if config.vocab_size is None:853 raise ValueError(854 f"You are trying to instantiate {self.__class__} with a configuration that "855 "does not define the vocabulary size of the language model head. Please "856 "instantiate the model as follows: `Data2VecAudioForCTC.from_pretrained(..., vocab_size=vocab_size)`. "857 "or define `vocab_size` of your model's configuration."858 )859 output_hidden_size = (860 config.output_hidden_size if hasattr(config, "add_adapter") and config.add_adapter else config.hidden_size861 )862 self.lm_head = nn.Linear(output_hidden_size, config.vocab_size)863 864 # Initialize weights and apply final processing865 self.post_init()866 867 def freeze_feature_extractor(self):868 """869 Calling this function will disable the gradient computation for the feature encoder so that its parameter will870 not be updated during training.871 """872 warnings.warn(873 "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "874 "Please use the equivalent `freeze_feature_encoder` method instead.",875 FutureWarning,876 )877 self.freeze_feature_encoder()878 879 def freeze_feature_encoder(self):880 """881 Calling this function will disable the gradient computation for the feature encoder so that its parameter will882 not be updated during training.883 """884 self.data2vec_audio.feature_extractor._freeze_parameters()885 886 @auto_docstring887 def forward(888 self,889 input_values: Optional[torch.Tensor],890 attention_mask: Optional[torch.Tensor] = None,891 output_attentions: Optional[bool] = None,892 output_hidden_states: Optional[bool] = None,893 return_dict: Optional[bool] = None,894 labels: Optional[torch.Tensor] = None,895 ) -> Union[tuple, CausalLMOutput]:896 r"""897 labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*):898 Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to899 the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`.900 All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,901 config.vocab_size - 1]`.902 """903 return_dict = return_dict if return_dict is not None else self.config.use_return_dict904 905 if labels is not None and labels.max() >= self.config.vocab_size:906 raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")907 908 outputs = self.data2vec_audio(909 input_values,910 attention_mask=attention_mask,911 output_attentions=output_attentions,912 output_hidden_states=output_hidden_states,913 return_dict=return_dict,914 )915 916 hidden_states = outputs[0]917 hidden_states = self.dropout(hidden_states)918 919 logits = self.lm_head(hidden_states)920 921 loss = None922 if labels is not None:923 # retrieve loss input_lengths from attention_mask924 attention_mask = (925 attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long)926 )927 input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)928 929 # assuming that padded tokens are filled with -100930 # when not being attended to931 labels_mask = labels >= 0932 target_lengths = labels_mask.sum(-1)933 flattened_targets = labels.masked_select(labels_mask)934 935 # ctc_loss doesn't support fp16936 log_probs = nn.functional.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)937 938 with torch.backends.cudnn.flags(enabled=False):939 loss = nn.functional.ctc_loss(940 log_probs,941 flattened_targets,942 input_lengths,943 target_lengths,944 blank=self.config.pad_token_id,945 reduction=self.config.ctc_loss_reduction,946 zero_infinity=self.config.ctc_zero_infinity,947 )948 949 if not return_dict:950 output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]951 return ((loss,) + output) if loss is not None else output952 953 return CausalLMOutput(954 loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions955 )956 957 958@auto_docstring(959 custom_intro="""960 Data2VecAudio Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like961 SUPERB Keyword Spotting.962 """963)964class Data2VecAudioForSequenceClassification(Data2VecAudioPreTrainedModel):965 def __init__(self, config):966 super().__init__(config)967 968 if hasattr(config, "add_adapter") and config.add_adapter:969 raise ValueError(970 "Sequence classification does not support the use of Data2VecAudio adapters (config.add_adapter=True)"971 )972 self.data2vec_audio = Data2VecAudioModel(config)973 num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings974 if config.use_weighted_layer_sum:975 self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)976 self.projector = nn.Linear(config.hidden_size, config.classifier_proj_size)977 self.classifier = nn.Linear(config.classifier_proj_size, config.num_labels)978 979 # Initialize weights and apply final processing980 self.post_init()981 982 def freeze_feature_extractor(self):983 """984 Calling this function will disable the gradient computation for the feature encoder so that its parameters will985 not be updated during training.986 """987 warnings.warn(988 "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "989 "Please use the equivalent `freeze_feature_encoder` method instead.",990 FutureWarning,991 )992 self.freeze_feature_encoder()993 994 def freeze_feature_encoder(self):995 """996 Calling this function will disable the gradient computation for the feature encoder so that its parameter will997 not be updated during training.998 """999 self.data2vec_audio.feature_extractor._freeze_parameters()1000 1001 def freeze_base_model(self):1002 """1003 Calling this function will disable the gradient computation for the base model so that its parameters will not1004 be updated during training. Only the classification head will be updated.1005 """1006 for param in self.data2vec_audio.parameters():1007 param.requires_grad = False1008 1009 @auto_docstring1010 def forward(1011 self,1012 input_values: Optional[torch.Tensor],1013 attention_mask: Optional[torch.Tensor] = None,1014 output_attentions: Optional[bool] = None,1015 output_hidden_states: Optional[bool] = None,1016 return_dict: Optional[bool] = None,1017 labels: Optional[torch.Tensor] = None,1018 ) -> Union[tuple, SequenceClassifierOutput]:1019 r"""1020 input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):1021 Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file1022 into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library1023 (`pip install torchcodec`) or the soundfile library (`pip install soundfile`).1024 To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion1025 into a tensor of type `torch.FloatTensor`. See [`Data2VecAudioProcessor.__call__`] for details.1026 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1027 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1028 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1029 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1030 """1031 1032 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1033 output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states1034 1035 outputs = self.data2vec_audio(1036 input_values,1037 attention_mask=attention_mask,1038 output_attentions=output_attentions,1039 output_hidden_states=output_hidden_states,1040 return_dict=return_dict,1041 )1042 1043 if self.config.use_weighted_layer_sum:1044 hidden_states = outputs[_HIDDEN_STATES_START_POSITION]1045 hidden_states = torch.stack(hidden_states, dim=1)1046 norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)1047 hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)1048 else:1049 hidden_states = outputs[0]1050 1051 hidden_states = self.projector(hidden_states)1052 if attention_mask is None:1053 pooled_output = hidden_states.mean(dim=1)1054 else:1055 padding_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)1056 expand_padding_mask = padding_mask.unsqueeze(-1).repeat(1, 1, hidden_states.shape[2])1057 hidden_states[~expand_padding_mask] = 0.01058 pooled_output = hidden_states.sum(dim=1) / padding_mask.sum(dim=1).view(-1, 1)1059 1060 logits = self.classifier(pooled_output)1061 1062 loss = None1063 if labels is not None:1064 loss_fct = CrossEntropyLoss()1065 loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))1066 1067 if not return_dict:1068 output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]1069 return ((loss,) + output) if loss is not None else output1070 1071 return SequenceClassifierOutput(1072 loss=loss,1073 logits=logits,1074 hidden_states=outputs.hidden_states,1075 attentions=outputs.attentions,1076 )1077 1078 1079@auto_docstring1080class Data2VecAudioForAudioFrameClassification(Data2VecAudioPreTrainedModel):1081 def __init__(self, config):1082 super().__init__(config)1083 1084 if hasattr(config, "add_adapter") and config.add_adapter:1085 raise ValueError(1086 "Audio frame classification does not support the use of Data2VecAudio adapters (config.add_adapter=True)"1087 )1088 self.data2vec_audio = Data2VecAudioModel(config)1089 num_layers = config.num_hidden_layers + 1 # transformer layers + input embeddings1090 if config.use_weighted_layer_sum:1091 self.layer_weights = nn.Parameter(torch.ones(num_layers) / num_layers)1092 self.classifier = nn.Linear(config.hidden_size, config.num_labels)1093 self.num_labels = config.num_labels1094 1095 self.init_weights()1096 1097 def freeze_feature_extractor(self):1098 """1099 Calling this function will disable the gradient computation for the feature encoder so that its parameter will1100 not be updated during training.1101 """1102 warnings.warn(1103 "The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. "1104 "Please use the equivalent `freeze_feature_encoder` method instead.",1105 FutureWarning,1106 )1107 self.freeze_feature_encoder()1108 1109 def freeze_feature_encoder(self):1110 """1111 Calling this function will disable the gradient computation for the feature encoder so that its parameter will1112 not be updated during training.1113 """1114 self.data2vec_audio.feature_extractor._freeze_parameters()1115 1116 def freeze_base_model(self):1117 """1118 Calling this function will disable the gradient computation for the base model so that its parameters will not1119 be updated during training. Only the classification head will be updated.1120 """1121 for param in self.data2vec_audio.parameters():1122 param.requires_grad = False1123 1124 @auto_docstring1125 def forward(1126 self,1127 input_values: Optional[torch.Tensor],1128 attention_mask: Optional[torch.Tensor] = None,1129 labels: Optional[torch.Tensor] = None,1130 output_attentions: Optional[bool] = None,1131 output_hidden_states: Optional[bool] = None,1132 return_dict: Optional[bool] = None,1133 ) -> Union[tuple, TokenClassifierOutput]:1134 r"""1135 input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):1136 Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file1137 into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library1138 (`pip install torchcodec`) or the soundfile library (`pip install soundfile`).1139 To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion1140 into a tensor of type `torch.FloatTensor`. See [`Data2VecAudioProcessor.__call__`] for details.1141 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1142 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1143 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1144 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1145 """1146 1147 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1148 output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states1149 1150 outputs = self.data2vec_audio(1151 input_values,1152 attention_mask=attention_mask,1153 output_attentions=output_attentions,1154 output_hidden_states=output_hidden_states,1155 return_dict=return_dict,1156 )1157 1158 if self.config.use_weighted_layer_sum:1159 hidden_states = outputs[_HIDDEN_STATES_START_POSITION]1160 hidden_states = torch.stack(hidden_states, dim=1)1161 norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)1162 hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)1163 else:1164 hidden_states = outputs[0]1165 1166 logits = self.classifier(hidden_states)1167 1168 loss = None1169 if labels is not None:1170 loss_fct = CrossEntropyLoss()1171 loss = loss_fct(logits.view(-1, self.num_labels), torch.argmax(labels.view(-1, self.num_labels), axis=1))1172 1173 if not return_dict:1174 output = (logits,) + outputs[_HIDDEN_STATES_START_POSITION:]1175 return output1176 1177 return TokenClassifierOutput(1178 loss=loss,1179 logits=logits,1180 hidden_states=outputs.hidden_states,1181 attentions=outputs.attentions,1182 )1183 1184 1185class AMSoftmaxLoss(nn.Module):1186 def __init__(self, input_dim, num_labels, scale=30.0, margin=0.4):1187 super().__init__()1188 self.scale = scale1189 self.margin = margin1190 self.num_labels = num_labels1191 self.weight = nn.Parameter(torch.randn(input_dim, num_labels), requires_grad=True)1192 self.loss = nn.CrossEntropyLoss()1193 1194 def forward(self, hidden_states, labels):1195 labels = labels.flatten()1196 weight = nn.functional.normalize(self.weight, dim=0)1197 hidden_states = nn.functional.normalize(hidden_states, dim=1)1198 cos_theta = torch.mm(hidden_states, weight)1199 psi = cos_theta - self.margin1200