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"""Flax Wav2Vec2 model."""16 17from functools import partial18from typing import Optional, Union19 20import flax21import flax.linen as nn22import jax23import jax.numpy as jnp24import numpy as np25from flax.core.frozen_dict import FrozenDict, freeze, unfreeze26from flax.linen.attention import dot_product_attention_weights27from flax.traverse_util import flatten_dict, unflatten_dict28from jax import lax29 30from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput31from ...modeling_flax_utils import (32 ACT2FN,33 FlaxPreTrainedModel,34 append_replace_return_docstrings,35 overwrite_call_docstring,36)37from ...utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging38from .configuration_wav2vec2 import Wav2Vec2Config39 40 41logger = logging.get_logger(__name__)42 43 44@flax.struct.dataclass45class FlaxWav2Vec2BaseModelOutput(ModelOutput):46 """47 Output type of [`FlaxWav2Vec2BaseModelOutput`], with potential hidden states and attentions.48 49 Args:50 last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`):51 Sequence of hidden-states at the output of the last layer of the model.52 extract_features (`jnp.ndarray` of shape `(batch_size, sequence_length, last_conv_dim)`):53 Sequence of extracted feature vectors of the last convolutional layer of the model with `last_conv_dim`54 being the dimension of the last convolutional layer.55 hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):56 Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape57 `(batch_size, sequence_length, hidden_size)`.58 59 Hidden-states of the model at the output of each layer plus the initial embedding outputs.60 attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):61 Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length,62 sequence_length)`.63 64 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention65 heads.66 """67 68 last_hidden_state: jnp.ndarray = None69 extract_features: jnp.ndarray = None70 hidden_states: Optional[tuple[jnp.ndarray]] = None71 attentions: Optional[tuple[jnp.ndarray]] = None72 73 74@flax.struct.dataclass75class FlaxWav2Vec2ForPreTrainingOutput(ModelOutput):76 """77 Output type of [`FlaxWav2Vec2ForPreTrainingOutput`], with potential hidden states and attentions.78 79 Args:80 loss (*optional*, returned when model is in train mode, `jnp.ndarray` of shape `(1,)`):81 Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official82 paper](https://huggingface.co/papers/2006.11477).83 projected_states (`jnp.ndarray` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):84 Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked85 projected quantized states.86 projected_quantized_states (`jnp.ndarray` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`):87 Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive88 target vectors for contrastive loss.89 hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):90 Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape91 `(batch_size, sequence_length, hidden_size)`.92 93 Hidden-states of the model at the output of each layer plus the initial embedding outputs.94 attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):95 Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length,96 sequence_length)`.97 98 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention99 heads.100 """101 102 projected_states: jnp.ndarray = None103 projected_quantized_states: jnp.ndarray = None104 codevector_perplexity: jnp.ndarray = None105 hidden_states: Optional[tuple[jnp.ndarray]] = None106 attentions: Optional[tuple[jnp.ndarray]] = None107 108 109def _compute_mask_indices(110 shape: tuple[int, int],111 mask_prob: float,112 mask_length: int,113 attention_mask: Optional[np.ndarray] = None,114 min_masks: int = 0,115) -> np.ndarray:116 """117 Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for118 ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on119 CPU as part of the preprocessing during training.120 121 Args:122 shape: the shape for which to compute masks.123 should be of size 2 where first element is batch size and 2nd is timesteps124 mask_prob:125 probability for each token to be chosen as start of the span to be masked. this will be multiplied by126 number of timesteps divided by length of mask span to mask approximately this percentage of all elements.127 however due to overlaps, the actual number will be smaller (unless no_overlap is True)128 mask_length: size of the mask129 min_masks: minimum number of masked spans130 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} and"140 f" `sequence_length`: {sequence_length}`"141 )142 143 # compute number of masked spans in batch144 num_masked_spans = int(mask_prob * sequence_length / mask_length + np.random.rand(1).item())145 num_masked_spans = max(num_masked_spans, min_masks)146 147 # make sure num masked indices <= sequence_length148 if num_masked_spans * mask_length > sequence_length:149 num_masked_spans = sequence_length // mask_length150 151 # SpecAugment mask to fill152 spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)153 154 # get random indices to mask155 spec_aug_mask_idxs = np.array(156 [157 np.random.choice(np.arange(sequence_length - (mask_length - 1)), num_masked_spans, replace=False)158 for _ in range(batch_size)159 ]160 )161 162 # expand masked indices to masked spans163 spec_aug_mask_idxs = np.broadcast_to(spec_aug_mask_idxs[:, :, None], (batch_size, num_masked_spans, mask_length))164 spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, num_masked_spans * mask_length)165 166 offsets = np.arange(mask_length)[None, None, :]167 offsets = np.broadcast_to(offsets, (batch_size, num_masked_spans, mask_length)).reshape(168 batch_size, num_masked_spans * mask_length169 )170 spec_aug_mask_idxs = spec_aug_mask_idxs + offsets171 172 # scatter indices to mask173 np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)174 175 if attention_mask is not None:176 # make sure padded input ids cannot be masked177 spec_aug_mask = np.where(attention_mask, spec_aug_mask, False)178 179 return spec_aug_mask180 181 182def _sample_negative_indices(features_shape: tuple, num_negatives: int, attention_mask: Optional[np.ndarray] = None):183 """184 Sample `num_negatives` vectors from feature vectors.185 """186 batch_size, sequence_length, hidden_size = features_shape187 if sequence_length <= 1:188 raise ValueError(189 "`features should have `sequence_length` > 1, but are of shape "190 f"(batch_size, sequence_length, hidden_size) = ({batch_size, sequence_length, hidden_size})."191 )192 193 # get `num_negatives` random vector indices from the same utterance194 sampled_negative_indices = []195 for batch_idx in range(batch_size):196 high = attention_mask[batch_idx].sum() - 1 if attention_mask is not None else sequence_length - 1197 sampled_indices_slice = np.random.randint(0, high, size=(num_negatives * sequence_length,))198 sampled_negative_indices.append(sampled_indices_slice)199 200 sampled_negative_indices = np.asarray(sampled_negative_indices, dtype=np.int32)201 202 # generate indices of the positive vectors themselves, repeat them `num_negatives` times203 feature_indices = np.broadcast_to(np.arange(sequence_length)[:, None], (sequence_length, num_negatives)).flatten()204 205 # avoid sampling the same positive vector, but keep the distribution uniform206 sampled_negative_indices[sampled_negative_indices >= feature_indices] += 1207 208 # correct for batch size209 for batch_idx in range(1, batch_size):210 sampled_negative_indices[batch_idx] += batch_idx * sequence_length211 212 return sampled_negative_indices213 214 215WAV2VEC2_START_DOCSTRING = r"""216 Wav2Vec2 was proposed in [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech217 Representations](https://huggingface.co/papers/2006.11477) by Alexei Baevski, Henry Zhou, Abdelrahman Mohamed, Michael218 Auli.219 220 This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the221 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads222 etc.)223 224 This model is also a Flax Linen225 [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a226 regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.227 228 Finally, this model supports inherent JAX features such as:229 230 - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)231 - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)232 - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)233 - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)234 235 Parameters:236 config ([`Wav2Vec2Config`]): Model configuration class with all the parameters of the model.237 Initializing with a config file does not load the weights associated with the model, only the238 configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.239 dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):240 The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and241 `jax.numpy.bfloat16` (on TPUs).242 243 This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If244 specified all the computation will be performed with the given `dtype`.245 246 **Note that this only specifies the dtype of the computation and does not influence the dtype of model247 parameters.**248 249 If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and250 [`~FlaxPreTrainedModel.to_bf16`].251"""252 253 254WAV2VEC2_INPUTS_DOCSTRING = r"""255 Args:256 input_values (`jnp.ndarray` of shape `(batch_size, sequence_length)`):257 Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file258 into an array of type `list[float]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library259 (`pip install torchcodec`) or the soundfile library (`pip install soundfile`).260 To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion261 into a tensor of type `jnp.ndarray`. See [`Wav2Vec2Processor.__call__`] for details.262 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):263 Mask to avoid performing convolution and attention on padding token indices. Mask values selected in `[0,264 1]`:265 266 - 1 for tokens that are **not masked**,267 - 0 for tokens that are **masked**.268 269 [What are attention masks?](../glossary#attention-mask) .. warning:: `attention_mask` should only be passed270 if the corresponding processor has `config.return_attention_mask == True`. For all models whose processor271 has `config.return_attention_mask == False`, such as272 [wav2vec2-base](https://huggingface.co/facebook/wav2vec2-base-960h), `attention_mask` should **not** be273 passed to avoid degraded performance when doing batched inference. For such models `input_values` should274 simply be padded with 0 and passed without `attention_mask`. Be aware that these models also yield slightly275 different results depending on whether `input_values` is padded or not.276 mask_time_indices (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):277 Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict278 masked extracted features in *config.proj_codevector_dim* space.279 output_attentions (`bool`, *optional*):280 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned281 tensors for more detail.282 output_hidden_states (`bool`, *optional*):283 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for284 more detail.285 return_dict (`bool`, *optional*):286 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.287"""288 289 290class FlaxWav2Vec2LayerNormConvLayer(nn.Module):291 config: Wav2Vec2Config292 layer_id: int = 0293 dtype: jnp.dtype = jnp.float32294 295 def setup(self):296 self.in_conv_dim = self.config.conv_dim[self.layer_id] if self.layer_id > 0 else 1297 self.out_conv_dim = self.config.conv_dim[self.layer_id]298 299 self.conv = nn.Conv(300 features=self.config.conv_dim[self.layer_id],301 kernel_size=(self.config.conv_kernel[self.layer_id],),302 strides=(self.config.conv_stride[self.layer_id],),303 use_bias=self.config.conv_bias,304 kernel_init=jax.nn.initializers.he_normal(),305 padding="VALID",306 dtype=self.dtype,307 )308 self.layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)309 self.activation = ACT2FN[self.config.feat_extract_activation]310 311 def __call__(self, hidden_states):312 hidden_states = self.conv(hidden_states)313 hidden_states = self.layer_norm(hidden_states)314 hidden_states = self.activation(hidden_states)315 return hidden_states316 317 318class FlaxConvWithWeightNorm(nn.Module):319 config: Wav2Vec2Config320 dtype: jnp.dtype = jnp.float32321 322 def setup(self):323 self.conv = nn.Conv(324 features=self.config.hidden_size,325 kernel_size=(self.config.num_conv_pos_embeddings,),326 kernel_init=jax.nn.initializers.he_normal(),327 padding="VALID",328 feature_group_count=self.config.num_conv_pos_embedding_groups,329 dtype=self.dtype,330 )331 weight_shape = (332 self.conv.features,333 self.conv.features // self.conv.feature_group_count,334 self.conv.kernel_size[0],335 )336 self.weight_v = self.param("weight_v", jax.nn.initializers.he_normal(), weight_shape)337 self.weight_g = self.param("weight_g", lambda _: jnp.linalg.norm(self.weight_v, axis=(0, 1))[None, None, :])338 self.bias = self.param("bias", jax.nn.initializers.zeros, (self.conv.features,))339 self.prev_padding = self.conv.kernel_size[0] // 2340 341 def _get_normed_weights(self):342 weight_v_norm = jnp.linalg.norm(self.weight_v, axis=(0, 1))[None, None, :]343 normed_weight_v = jnp.divide(self.weight_v, weight_v_norm)344 normed_kernel = jnp.multiply(normed_weight_v, self.weight_g)345 return normed_kernel346 347 def __call__(self, hidden_states):348 kernel = self._get_normed_weights()349 hidden_states = jnp.pad(hidden_states, ((0, 0), (self.prev_padding, self.prev_padding), (0, 0)))350 hidden_states = self.conv.apply({"params": {"kernel": kernel.T, "bias": self.bias}}, hidden_states)351 return hidden_states352 353 354class FlaxWav2Vec2PositionalConvEmbedding(nn.Module):355 config: Wav2Vec2Config356 dtype: jnp.dtype = jnp.float32357 358 def setup(self):359 self.conv = FlaxConvWithWeightNorm(self.config, dtype=self.dtype)360 self.activation = ACT2FN[self.config.feat_extract_activation]361 self.num_pad_remove = 1 if self.config.num_conv_pos_embeddings % 2 == 0 else 0362 363 def __call__(self, hidden_states):364 hidden_states = hidden_states.transpose((0, 1, 2))365 366 hidden_states = self.conv(hidden_states)367 368 if self.num_pad_remove > 0:369 hidden_states = hidden_states[:, : -self.num_pad_remove, :]370 hidden_states = self.activation(hidden_states)371 372 hidden_states = hidden_states.transpose((0, 1, 2))373 return hidden_states374 375 376class FlaxConvLayersCollection(nn.Module):377 config: Wav2Vec2Config378 dtype: jnp.dtype = jnp.float32379 380 def setup(self):381 if self.config.feat_extract_norm == "layer":382 self.layers = [383 FlaxWav2Vec2LayerNormConvLayer(self.config, layer_id=i, name=str(i), dtype=self.dtype)384 for i in range(self.config.num_feat_extract_layers)385 ]386 elif self.config.feat_extract_norm == "group":387 raise NotImplementedError("At the moment only ``config.feat_extract_norm == 'layer'`` is supported")388 else:389 raise ValueError(390 f"`config.feat_extract_norm` is {self.config.feat_extract_norm}, but has to be one of ['group',"391 " 'layer']"392 )393 394 def __call__(self, hidden_states):395 for i, conv_layer in enumerate(self.layers):396 hidden_states = conv_layer(hidden_states)397 return hidden_states398 399 400class FlaxWav2Vec2FeatureEncoder(nn.Module):401 """Construct the features from raw audio waveform"""402 403 config: Wav2Vec2Config404 dtype: jnp.dtype = jnp.float32405 406 def setup(self):407 self.conv_layers = FlaxConvLayersCollection(self.config, dtype=self.dtype)408 409 def __call__(self, input_values, freeze_feature_encoder=False):410 hidden_states = input_values[:, :, None]411 hidden_states = self.conv_layers(hidden_states)412 if freeze_feature_encoder:413 hidden_states = jax.lax.stop_gradient(hidden_states)414 return hidden_states415 416 417class FlaxWav2Vec2FeatureProjection(nn.Module):418 config: Wav2Vec2Config419 dtype: jnp.dtype = jnp.float32420 421 def setup(self):422 self.layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)423 self.projection = nn.Dense(424 self.config.hidden_size,425 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),426 dtype=self.dtype,427 )428 self.dropout = nn.Dropout(rate=self.config.feat_proj_dropout)429 430 def __call__(self, hidden_states, deterministic=True):431 norm_hidden_states = self.layer_norm(hidden_states)432 hidden_states = self.projection(norm_hidden_states)433 hidden_states = self.dropout(hidden_states, deterministic=deterministic)434 return hidden_states, norm_hidden_states435 436 437class FlaxWav2Vec2Attention(nn.Module):438 config: Wav2Vec2Config439 embed_dim: int440 num_heads: int441 dropout: float = 0.0442 bias: bool = True443 dtype: jnp.dtype = jnp.float32 # the dtype of the computation444 445 def setup(self) -> None:446 self.head_dim = self.embed_dim // self.num_heads447 if self.head_dim * self.num_heads != self.embed_dim:448 raise ValueError(449 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"450 f" {self.num_heads})."451 )452 453 dense = partial(454 nn.Dense,455 self.embed_dim,456 use_bias=self.bias,457 dtype=self.dtype,458 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),459 )460 461 self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()462 self.out_proj = dense()463 464 self.dropout_layer = nn.Dropout(rate=self.dropout)465 466 def _split_heads(self, hidden_states):467 return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))468 469 def _merge_heads(self, hidden_states):470 return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))471 472 def __call__(473 self,474 hidden_states: jnp.ndarray,475 key_value_states: Optional[jnp.ndarray] = None,476 attention_mask: Optional[jnp.ndarray] = None,477 deterministic: bool = True,478 ) -> tuple[jnp.ndarray]:479 """Input shape: Batch x Time x Channel"""480 481 # get query proj482 query_states = self.q_proj(hidden_states)483 484 key_states = self.k_proj(hidden_states)485 value_states = self.v_proj(hidden_states)486 487 query_states = self._split_heads(query_states)488 key_states = self._split_heads(key_states)489 value_states = self._split_heads(value_states)490 491 if attention_mask is not None:492 attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))493 494 # Convert the boolean attention mask to an attention bias.495 if attention_mask is not None:496 # attention mask in the form of attention bias497 attention_bias = lax.select(498 attention_mask > 0,499 jnp.full(attention_mask.shape, 0.0).astype(self.dtype),500 jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),501 )502 else:503 attention_bias = None504 505 dropout_rng = None506 if not deterministic and self.dropout > 0.0:507 dropout_rng = self.make_rng("dropout")508 509 attn_weights = dot_product_attention_weights(510 query_states,511 key_states,512 bias=attention_bias,513 dropout_rng=dropout_rng,514 dropout_rate=self.dropout,515 broadcast_dropout=True,516 deterministic=deterministic,517 dtype=self.dtype,518 precision=None,519 )520 521 attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)522 attn_output = self._merge_heads(attn_output)523 attn_output = self.out_proj(attn_output)524 525 return attn_output, attn_weights526 527 528class FlaxWav2Vec2FeedForward(nn.Module):529 config: Wav2Vec2Config530 dtype: jnp.dtype = jnp.float32531 532 def setup(self):533 self.intermediate_dropout = nn.Dropout(rate=self.config.activation_dropout)534 535 self.intermediate_dense = nn.Dense(536 self.config.intermediate_size,537 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),538 dtype=self.dtype,539 )540 if isinstance(self.config.hidden_act, str):541 self.intermediate_act_fn = ACT2FN[self.config.hidden_act]542 else:543 self.intermediate_act_fn = self.config.hidden_act544 545 self.output_dense = nn.Dense(546 self.config.hidden_size,547 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),548 dtype=self.dtype,549 )550 self.output_dropout = nn.Dropout(rate=self.config.hidden_dropout)551 552 def __call__(self, hidden_states, deterministic=True):553 hidden_states = self.intermediate_dense(hidden_states)554 hidden_states = self.intermediate_act_fn(hidden_states)555 hidden_states = self.intermediate_dropout(hidden_states, deterministic=deterministic)556 557 hidden_states = self.output_dense(hidden_states)558 hidden_states = self.output_dropout(hidden_states, deterministic=deterministic)559 return hidden_states560 561 562class FlaxWav2Vec2EncoderLayerStableLayerNorm(nn.Module):563 config: Wav2Vec2Config564 dtype: jnp.dtype = jnp.float32565 566 def setup(self):567 self.attention = FlaxWav2Vec2Attention(568 config=self.config,569 embed_dim=self.config.hidden_size,570 num_heads=self.config.num_attention_heads,571 dropout=self.config.attention_dropout,572 dtype=self.dtype,573 )574 self.dropout = nn.Dropout(rate=self.config.hidden_dropout)575 self.layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)576 self.feed_forward = FlaxWav2Vec2FeedForward(self.config, dtype=self.dtype)577 self.final_layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)578 579 def __call__(self, hidden_states, attention_mask=None, deterministic=True, output_attentions=False):580 attn_residual = hidden_states581 hidden_states = self.layer_norm(hidden_states)582 hidden_states, attn_weights = self.attention(583 hidden_states, attention_mask=attention_mask, deterministic=deterministic584 )585 hidden_states = self.dropout(hidden_states, deterministic=deterministic)586 hidden_states = attn_residual + hidden_states587 hidden_states = hidden_states + self.feed_forward(588 self.final_layer_norm(hidden_states), deterministic=deterministic589 )590 591 outputs = (hidden_states,)592 593 if output_attentions:594 outputs += (attn_weights,)595 596 return outputs597 598 599class FlaxWav2Vec2EncoderLayerStableLayerNormCollection(nn.Module):600 config: Wav2Vec2Config601 dtype: jnp.dtype = jnp.float32602 603 def setup(self):604 self.layers = [605 FlaxWav2Vec2EncoderLayerStableLayerNorm(self.config, name=str(i), dtype=self.dtype)606 for i in range(self.config.num_hidden_layers)607 ]608 609 def __call__(610 self,611 hidden_states,612 attention_mask=None,613 deterministic: bool = True,614 output_attentions: bool = False,615 output_hidden_states: bool = False,616 return_dict: bool = True,617 ):618 all_attentions = () if output_attentions else None619 all_hidden_states = () if output_hidden_states else None620 621 for i, layer in enumerate(self.layers):622 if output_hidden_states:623 all_hidden_states += (hidden_states,)624 625 layer_outputs = layer(626 hidden_states, attention_mask, deterministic=deterministic, output_attentions=output_attentions627 )628 629 hidden_states = layer_outputs[0]630 631 if output_attentions:632 all_attentions += (layer_outputs[1],)633 634 if output_hidden_states:635 all_hidden_states += (hidden_states,)636 637 outputs = (hidden_states, all_hidden_states, all_attentions)638 639 if not return_dict:640 return tuple(v for v in outputs if v is not None)641 642 return FlaxBaseModelOutput(643 last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions644 )645 646 647class FlaxWav2Vec2StableLayerNormEncoder(nn.Module):648 config: Wav2Vec2Config649 dtype: jnp.dtype = jnp.float32650 651 def setup(self):652 self.pos_conv_embed = FlaxWav2Vec2PositionalConvEmbedding(self.config, dtype=self.dtype)653 self.layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)654 self.dropout = nn.Dropout(rate=self.config.hidden_dropout)655 self.layers = FlaxWav2Vec2EncoderLayerStableLayerNormCollection(self.config, dtype=self.dtype)656 657 def __call__(658 self,659 hidden_states,660 attention_mask=None,661 deterministic=True,662 output_attentions=False,663 output_hidden_states=False,664 return_dict=True,665 ):666 if attention_mask is not None:667 # make sure padded tokens are not attended to668 hidden_states = jnp.where(669 jnp.broadcast_to(attention_mask[:, :, None], hidden_states.shape), hidden_states, 0670 )671 672 position_embeddings = self.pos_conv_embed(hidden_states)673 674 hidden_states = hidden_states + position_embeddings675 hidden_states = self.dropout(hidden_states, deterministic=deterministic)676 677 outputs = self.layers(678 hidden_states,679 attention_mask,680 output_attentions=output_attentions,681 output_hidden_states=output_hidden_states,682 return_dict=return_dict,683 )684 685 last_hidden_state = self.layer_norm(outputs[0])686 687 # update the last element in `hidden_states` after applying `layernorm` above688 hidden_states = None689 if output_hidden_states:690 hidden_states = outputs[1]691 hidden_states = hidden_states[:-1] + (last_hidden_state,)692 693 if not return_dict:694 outputs = (last_hidden_state, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])695 return tuple(v for v in outputs if v is not None)696 697 return FlaxBaseModelOutput(698 last_hidden_state=last_hidden_state, hidden_states=hidden_states, attentions=outputs.attentions699 )700 701 702class FlaxWav2Vec2GumbelVectorQuantizer(nn.Module):703 """704 Vector quantization using gumbel softmax. See [CATEGORICAL REPARAMETERIZATION WITH705 GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) for more information.706 """707 708 config: Wav2Vec2Config709 dtype: jnp.dtype = jnp.float32710 711 def setup(self):712 self.num_groups = self.config.num_codevector_groups713 self.num_vars = self.config.num_codevectors_per_group714 715 if self.config.codevector_dim % self.num_groups != 0:716 raise ValueError(717 f"`config.codevector_dim {self.config.codevector_dim} must be divisible by"718 f" `config.num_codevector_groups` {self.num_groups} for concatenation"719 )720 721 # storage for codebook variables (codewords)722 self.codevectors = self.param(723 "codevectors",724 jax.nn.initializers.uniform(),725 (1, self.num_groups * self.num_vars, self.config.codevector_dim // self.num_groups),726 )727 self.weight_proj = nn.Dense(728 self.num_groups * self.num_vars,729 kernel_init=jax.nn.initializers.normal(1.0),730 dtype=self.dtype,731 )732 733 @staticmethod734 def _compute_perplexity(probs, mask=None):735 if mask is not None:736 mask_extended = jnp.broadcast_to(mask.flatten()[:, None, None], probs.shape)737 probs = jnp.where(mask_extended, probs, jnp.zeros_like(probs))738 marginal_probs = probs.sum(axis=0) / mask.sum()739 else:740 marginal_probs = probs.mean(axis=0)741 742 perplexity = jnp.exp(-jnp.sum(marginal_probs * jnp.log(marginal_probs + 1e-7), axis=-1)).sum()743 return perplexity744 745 def __call__(self, hidden_states, mask_time_indices=None, deterministic=True, temperature=1):746 batch_size, sequence_length, hidden_size = hidden_states.shape747 748 # project to codevector dim749 hidden_states = self.weight_proj(hidden_states)750 hidden_states = hidden_states.reshape(batch_size * sequence_length * self.num_groups, -1)751 752 if not deterministic:753 # sample code vector probs via gumbel in differentiateable way754 gumbel_rng = self.make_rng("gumbel")755 gumbels = jax.random.gumbel(gumbel_rng, hidden_states.shape)756 codevector_probs = nn.softmax((hidden_states + gumbels) / temperature)757 758 # compute perplexity759 codevector_soft_dist = nn.softmax(760 hidden_states.reshape(batch_size * sequence_length, self.num_groups, -1), axis=-1761 )762 perplexity = self._compute_perplexity(codevector_soft_dist, mask_time_indices)763 else:764 # take argmax in non-differentiable way765 # comptute hard codevector distribution (one hot)766 codevector_idx = hidden_states.argmax(axis=-1)767 codevector_probs = jax.nn.one_hot(codevector_idx, hidden_states.shape[-1]) * 1.0768 codevector_probs = codevector_probs.reshape(batch_size * sequence_length, self.num_groups, -1)769 perplexity = self._compute_perplexity(codevector_probs, mask_time_indices)770 771 codevector_probs = codevector_probs.reshape(batch_size * sequence_length, -1)772 # use probs to retrieve codevectors773 codevectors_per_group = jnp.expand_dims(codevector_probs, axis=-1) * self.codevectors774 codevectors = codevectors_per_group.reshape(batch_size * sequence_length, self.num_groups, self.num_vars, -1)775 codevectors = codevectors.sum(-2).reshape(batch_size, sequence_length, -1)776 777 return codevectors, perplexity778 779 780class FlaxWav2Vec2Adapter(nn.Module):781 config: Wav2Vec2Config782 dtype: jnp.dtype = jnp.float32783 784 def setup(self):785 # hidden_states require down-projection if feature dims don't match786 if self.config.output_hidden_size != self.config.hidden_size:787 self.proj = nn.Dense(788 self.config.output_hidden_size,789 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),790 dtype=self.dtype,791 )792 self.proj_layer_norm = nn.LayerNorm(epsilon=self.config.layer_norm_eps, dtype=self.dtype)793 else:794 self.proj = self.proj_layer_norm = None795 796 self.layers = FlaxWav2Vec2AdapterLayersCollection(self.config, dtype=self.dtype)797 798 def __call__(self, hidden_states, deterministic=True):799 # down-project hidden_states if required800 if self.proj is not None and self.proj_layer_norm is not None:801 hidden_states = self.proj(hidden_states)802 hidden_states = self.proj_layer_norm(hidden_states)803 804 hidden_states = self.layers(hidden_states)805 806 return hidden_states807 808 809class FlaxWav2Vec2AdapterLayer(nn.Module):810 config: Wav2Vec2Config811 dtype: jnp.dtype = jnp.float32812 813 def setup(self):814 self.conv = nn.Conv(815 features=2 * self.config.output_hidden_size,816 kernel_size=(self.config.adapter_kernel_size,),817 strides=(self.config.adapter_stride,),818 padding=((1, 1),),819 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),820 dtype=self.dtype,821 )822 823 def __call__(self, hidden_states):824 hidden_states = self.conv(hidden_states)825 hidden_states = nn.glu(hidden_states, axis=2)826 827 return hidden_states828 829 830class FlaxWav2Vec2AdapterLayersCollection(nn.Module):831 config: Wav2Vec2Config832 dtype: jnp.dtype = jnp.float32833 834 def setup(self):835 self.layers = [836 FlaxWav2Vec2AdapterLayer(self.config, name=str(i), dtype=self.dtype)837 for i in range(self.config.num_adapter_layers)838 ]839 840 def __call__(self, hidden_states):841 for conv_layer in self.layers:842 hidden_states = conv_layer(hidden_states)843 844 return hidden_states845 846 847class FlaxWav2Vec2PreTrainedModel(FlaxPreTrainedModel):848 """849 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained850 models.851 """852 853 config_class = Wav2Vec2Config854 base_model_prefix: str = "wav2vec2"855 main_input_name = "input_values"856 module_class: nn.Module = None857 858 def __init__(859 self,860 config: Wav2Vec2Config,861 input_shape: tuple = (1, 1024),862 seed: int = 0,863 dtype: jnp.dtype = jnp.float32,864 _do_init: bool = True,865 **kwargs,866 ):867 module = self.module_class(config=config, dtype=dtype, **kwargs)868 super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)869 870 def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:871 # init input tensors872 input_values = jnp.zeros(input_shape, dtype="i4")873 attention_mask = jnp.ones_like(input_values)874 params_rng, dropout_rng = jax.random.split(rng, 2)875 rngs = {"params": params_rng, "dropout": dropout_rng}876 877 random_params = self.module.init(rngs, input_values, attention_mask, return_dict=False)["params"]878 879 if params is not None:880 random_params = flatten_dict(unfreeze(random_params))881 params = flatten_dict(unfreeze(params))882 for missing_key in self._missing_keys:883 params[missing_key] = random_params[missing_key]884 self._missing_keys = set()885 return freeze(unflatten_dict(params))886 else:887 return random_params888 889 @add_start_docstrings_to_model_forward(WAV2VEC2_INPUTS_DOCSTRING)890 def __call__(891 self,892 input_values,893 attention_mask=None,894 mask_time_indices=None,895 params: Optional[dict] = None,896 dropout_rng: jax.random.PRNGKey = None,897 train: bool = False,898 output_attentions: Optional[bool] = None,899 output_hidden_states: Optional[bool] = None,900 freeze_feature_encoder: bool = False,901 return_dict: Optional[bool] = None,902 ):903 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions904 output_hidden_states = (905 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states906 )907 return_dict = return_dict if return_dict is not None else self.config.return_dict908 909 batch_size, sequence_length = input_values.shape910 911 if attention_mask is None:912 attention_mask = jnp.ones((batch_size, sequence_length))913 914 # Handle any PRNG if needed915 rngs = {}916 if dropout_rng is not None:917 rngs["dropout"] = dropout_rng918 919 inputs = {"params": params or self.params}920 921 return self.module.apply(922 inputs,923 jnp.array(input_values, dtype="f4"),924 jnp.array(attention_mask, dtype="i4"),925 mask_time_indices,926 not train,927 output_attentions,928 output_hidden_states,929 freeze_feature_encoder,930 return_dict,931 rngs=rngs,932 )933 934 def _get_feat_extract_output_lengths(935 self, input_lengths: Union[jnp.ndarray, int], add_adapter: Optional[bool] = None936 ):937 return self.module._get_feat_extract_output_lengths(input_lengths, add_adapter=add_adapter)938 939 940class FlaxWav2Vec2Module(nn.Module):941 config: Wav2Vec2Config942 dtype: jnp.dtype = jnp.float32943 944 def setup(self):945 self.feature_extractor = FlaxWav2Vec2FeatureEncoder(self.config, dtype=self.dtype)946 self.feature_projection = FlaxWav2Vec2FeatureProjection(self.config, dtype=self.dtype)947 self.masked_spec_embed = self.param(948 "masked_spec_embed", jax.nn.initializers.uniform(), (self.config.hidden_size,)949 )950 951 if self.config.do_stable_layer_norm:952 self.encoder = FlaxWav2Vec2StableLayerNormEncoder(self.config, dtype=self.dtype)953 else:954 raise NotImplementedError("``config.do_stable_layer_norm is False`` is currently not supported.")955 956 self.adapter = FlaxWav2Vec2Adapter(self.config, dtype=self.dtype) if self.config.add_adapter else None957 958 def __call__(959 self,960 input_values,961 attention_mask=None,962 mask_time_indices=None,963 deterministic=True,964 output_attentions=None,965 output_hidden_states=None,966 freeze_feature_encoder=False,967 return_dict=None,968 ):969 extract_features = self.feature_extractor(input_values, freeze_feature_encoder=freeze_feature_encoder)970 971 # make sure that no loss is computed on padded inputs972 if attention_mask is not None:973 # compute reduced attention_mask corresponding to feature vectors974 attention_mask = self._get_feature_vector_attention_mask(975 extract_features.shape[1], attention_mask, add_adapter=False976 )977 978 hidden_states, extract_features = self.feature_projection(extract_features, deterministic=deterministic)979 if mask_time_indices is not None: # apply SpecAugment along time axis with given indices980 hidden_states = jnp.where(981 jnp.broadcast_to(mask_time_indices[:, :, None], hidden_states.shape),982 jnp.broadcast_to(self.masked_spec_embed[None, None, :], hidden_states.shape),983 hidden_states,984 )985 986 encoder_outputs = self.encoder(987 hidden_states,988 attention_mask=attention_mask,989 deterministic=deterministic,990 output_attentions=output_attentions,991 output_hidden_states=output_hidden_states,992 return_dict=return_dict,993 )994 995 hidden_states = encoder_outputs[0]996 997 if self.adapter is not None:998 hidden_states = self.adapter(hidden_states)999 1000 if not return_dict:1001 return (hidden_states, extract_features) + encoder_outputs[1:]1002 1003 return FlaxWav2Vec2BaseModelOutput(1004 last_hidden_state=hidden_states,1005 extract_features=extract_features,1006 hidden_states=encoder_outputs.hidden_states,1007 attentions=encoder_outputs.attentions,1008 )1009 1010 def _get_feat_extract_output_lengths(1011 self, input_lengths: Union[jnp.ndarray, int], add_adapter: Optional[bool] = None1012 ):1013 """1014 Computes the output length of the convolutional layers1015 """1016 1017 add_adapter = self.config.add_adapter if add_adapter is None else add_adapter1018 1019 def _conv_out_length(input_length, kernel_size, stride):1020 # 1D convolutional layer output length formula taken1021 # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html1022 return (input_length - kernel_size) // stride + 11023 1024 for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):1025 input_lengths = _conv_out_length(input_lengths, kernel_size, stride)1026 1027 if add_adapter:1028 for _ in range(self.config.num_adapter_layers):1029 input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)1030 1031 return input_lengths1032 1033 def _get_feature_vector_attention_mask(1034 self, feature_vector_length: int, attention_mask: jnp.ndarray, add_adapter=None1035 ):1036 # Effectively attention_mask.sum(-1), but not inplace to be able to run1037 # on inference mode.1038 non_padded_lengths = attention_mask.cumsum(axis=-1)[:, -1]1039 1040 output_lengths = self._get_feat_extract_output_lengths(non_padded_lengths, add_adapter=add_adapter)1041 1042 batch_size = attention_mask.shape[0]1043 1044 attention_mask = jnp.zeros((batch_size, feature_vector_length), dtype=attention_mask.dtype)1045 # these two operations makes sure that all values1046 # before the output lengths indices are attended to1047 attention_mask = attention_mask.at[jnp.arange(attention_mask.shape[0]), output_lengths - 1].set(1)1048 attention_mask = jnp.flip(jnp.flip(attention_mask, -1).cumsum(-1), -1).astype("bool")1049 return attention_mask1050 1051 1052@add_start_docstrings(1053 "The bare Wav2Vec2 Model transformer outputting raw hidden-states without any specific head on top.",1054 WAV2VEC2_START_DOCSTRING,1055)1056class FlaxWav2Vec2Model(FlaxWav2Vec2PreTrainedModel):1057 module_class = FlaxWav2Vec2Module1058 1059 1060FLAX_WAV2VEC2_MODEL_DOCSTRING = """1061 Returns:1062 1063 Example:1064 1065 ```python1066 >>> from transformers import AutoProcessor, FlaxWav2Vec2Model1067 >>> from datasets import load_dataset1068 1069 >>> processor = AutoProcessor.from_pretrained("facebook/wav2vec2-large-lv60")1070 >>> model = FlaxWav2Vec2Model.from_pretrained("facebook/wav2vec2-large-lv60")1071 1072 1073 >>> def map_to_array(example):1074 ... example["speech"] = example["audio"]["array"]1075 ... return example1076 1077 1078 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")1079 >>> ds = ds.map(map_to_array)1080 1081 >>> input_values = processor(1082 ... ds["speech"][0], sampling_rate=16_000, return_tensors="np"1083 ... ).input_values # Batch size 11084 >>> hidden_states = model(input_values).last_hidden_state1085 ```1086"""1087 1088overwrite_call_docstring(1089 FlaxWav2Vec2Model,1090 WAV2VEC2_INPUTS_DOCSTRING + FLAX_WAV2VEC2_MODEL_DOCSTRING,1091)1092append_replace_return_docstrings(1093 FlaxWav2Vec2Model, output_type=FlaxWav2Vec2BaseModelOutput, config_class=Wav2Vec2Config1094)1095 1096 1097class FlaxWav2Vec2ForCTCModule(nn.Module):1098 config: Wav2Vec2Config1099 dtype: jnp.dtype = jnp.float321100 1101 def setup(self):1102 self.wav2vec2 = FlaxWav2Vec2Module(self.config, dtype=self.dtype)1103 self.dropout = nn.Dropout(rate=self.config.final_dropout)1104 self.lm_head = nn.Dense(1105 self.config.vocab_size,1106 kernel_init=jax.nn.initializers.normal(self.config.initializer_range),1107 dtype=self.dtype,1108 )1109 1110 def __call__(1111 self,1112 input_values,1113 attention_mask=None,1114 mask_time_indices=None,1115 deterministic=True,1116 output_attentions=None,1117 output_hidden_states=None,1118 freeze_feature_encoder=False,1119 return_dict=None,1120 ):1121 outputs = self.wav2vec2(1122 input_values,1123 attention_mask=attention_mask,1124 mask_time_indices=mask_time_indices,1125 deterministic=deterministic,1126 output_attentions=output_attentions,1127 output_hidden_states=output_hidden_states,1128 freeze_feature_encoder=freeze_feature_encoder,1129 return_dict=return_dict,1130 )1131 1132 hidden_states = outputs[0]1133 hidden_states = self.dropout(hidden_states, deterministic=deterministic)1134 1135 logits = self.lm_head(hidden_states)1136 1137 if not return_dict:1138 return (logits,) + outputs[2:]1139 1140 return FlaxCausalLMOutput(logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)1141 1142 def _get_feat_extract_output_lengths(1143 self,1144 input_lengths: Union[jnp.ndarray, int],1145 add_adapter: Optional[bool] = None,1146 ):1147 """1148 Computes the output length of the convolutional layers1149 """1150 1151 add_adapter = self.config.add_adapter if add_adapter is None else add_adapter1152 1153 def _conv_out_length(input_length, kernel_size, stride):1154 # 1D convolutional layer output length formula taken1155 # from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html1156 return (input_length - kernel_size) // stride + 11157 1158 for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride):1159 input_lengths = _conv_out_length(input_lengths, kernel_size, stride)1160 1161 if add_adapter:1162 for _ in range(self.config.num_adapter_layers):1163 input_lengths = _conv_out_length(input_lengths, 1, self.config.adapter_stride)1164 1165 return input_lengths1166 1167 1168@add_start_docstrings(1169 "Wav2Vec2 Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).",1170 WAV2VEC2_START_DOCSTRING,1171)1172class FlaxWav2Vec2ForCTC(FlaxWav2Vec2PreTrainedModel):1173 module_class = FlaxWav2Vec2ForCTCModule1174 1175 1176FLAX_WAV2VEC2_FOR_CTC_DOCSTRING = """1177 Returns:1178 1179 Example:1180 1181 ```python1182 >>> import jax.numpy as jnp1183 >>> from transformers import AutoProcessor, FlaxWav2Vec2ForCTC1184 >>> from datasets import load_dataset1185 1186 >>> processor = AutoProcessor.from_pretrained("facebook/wav2vec2-large-960h-lv60")1187 >>> model = FlaxWav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h-lv60")1188 1189 1190 >>> def map_to_array(example):1191 ... example["speech"] = example["audio"]["array"]1192 ... return example1193 1194 1195 >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")1196 >>> ds = ds.map(map_to_array)1197 1198 >>> input_values = processor(1199 ... ds["speech"][0], sampling_rate=16_000, return_tensors="np"1200 ... ).input_values # Batch size 1