# Copyright (c) Facebook, Inc. and its affiliates. # Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Paddle Wav2Vec2 model.""" import math import uuid from dataclasses import dataclass from dataclasses import field from enum import Enum from enum import EnumMeta from typing import Callable from typing import Dict from typing import List from typing import Optional from typing import Tuple import numpy as np import paddle import paddle.nn as nn import paddle.nn.functional as F from paddle import Tensor from paddlespeech.s2t.modules.align import Conv1D from paddlespeech.s2t.modules.align import Conv2D from paddlespeech.s2t.modules.align import Embedding from paddlespeech.s2t.modules.align import LayerNorm from paddlespeech.s2t.modules.align import Linear from paddlespeech.s2t.utils.log import Log logger = Log(__name__).getlog() class GLU(nn.Layer): r"""Applies the gated linear unit function :math:`{GLU}(a, b)= a \otimes \sigma(b)` where :math:`a` is the first half of the input matrices and :math:`b` is the second half. Args: axis (int): the dimension on which to split the input. Default: -1 Shape: - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional dimensions - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2` Examples:: >>> m = nn.GLU() >>> input = paddle.randn([4, 2]) >>> output = m(input) """ def __init__(self, axis: int=-1) -> None: super().__init__() self.axis = axis def forward(self, input: Tensor) -> Tensor: return F.glu(input, self.axis) class FairseqIncrementalState(object): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.init_incremental_state() def init_incremental_state(self): self._incremental_state_id = str(uuid.uuid4()) def _get_full_incremental_state_key(self, key: str) -> str: return "{}.{}".format(self._incremental_state_id, key) def get_incremental_state( self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]], key: str, ) -> Optional[Dict[str, Optional[Tensor]]]: """Helper for getting incremental state for an nn.Layer.""" full_key = self._get_full_incremental_state_key(key) if incremental_state is None or full_key not in incremental_state: return None return incremental_state[full_key] def set_incremental_state( self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]], key: str, value: Dict[str, Optional[Tensor]], ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]: """Helper for setting incremental state for an nn.Layer.""" if incremental_state is not None: full_key = self._get_full_incremental_state_key(key) incremental_state[full_key] = value return incremental_state def with_incremental_state(cls): cls.__bases__ = (FairseqIncrementalState, ) + tuple( b for b in cls.__bases__ if b != FairseqIncrementalState) return cls class FairseqDropout(paddle.nn.Layer): def __init__(self, p, module_name=None): super().__init__() self.p = p self.module_name = module_name self.apply_during_inference = False def forward(self, x): if self.p > 0 and (self.training or self.apply_during_inference): return F.dropout(x, p=self.p, training=True) else: return x def make_generation_fast_( self, name: str, retain_dropout: bool=False, retain_dropout_modules: Optional[List[str]]=None, **kwargs, ): if retain_dropout: if retain_dropout_modules is not None and self.module_name is None: logger.warning( "Cannot enable dropout during inference for module {} " "because module_name was not set".format(name)) elif (retain_dropout_modules is None # if None, apply to all modules or self.module_name in retain_dropout_modules): logger.info("Enabling dropout during inference for module: {}". format(name)) self.apply_during_inference = True else: logger.info("Disabling dropout for module: {}".format(name)) def quant_noise(module, p, block_size): """ Wraps modules and applies quantization noise to the weights for subsequent quantization with Iterative Product Quantization as described in "Training with Quantization Noise for Extreme Model Compression" Args: - module: nn.Layer - p: amount of Quantization Noise - block_size: size of the blocks for subsequent quantization with iPQ Remarks: - Layer weights must have the right sizes wrt the block size - Only Linear, Embedding and Conv2d modules are supported for the moment - For more detail on how to quantize by blocks with convolutional weights, see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks" - We implement the simplest form of noise here as stated in the paper which consists in randomly dropping blocks """ # if no quantization noise, don't register hook if p <= 0: return module # supported modules assert isinstance(module, (Linear, Embedding, Conv2D)) # test whether module.weight has the right sizes wrt block_size is_conv = len(module.weight.shape) == 4 # 2D matrix if not is_conv: if isinstance(module, Linear): features_weight = module.weight.shape[0] else: features_weight = module.weight.shape[1] assert ( features_weight % block_size == 0), "Input features must be a multiple of block sizes" # 4D matrix else: # 1x1 convolutions if module.weight.shape[2:] == (1, 1): assert (module.weight.shape[1] % block_size == 0 ), "Input channels must be a multiple of block sizes" # regular convolutions else: k = module.weight.shape[2] * module.weight.shape[3] assert k % block_size == 0, "Kernel size must be a multiple of block size" def _forward_pre_hook(mod, input): # no noise for evaluation if mod.training: if not is_conv: # gather weight and sizes weight = mod.weight if isinstance(module, Linear): in_features = weight.shape[0] out_features = weight.shape[1] else: in_features = weight.shape[1] out_features = weight.shape[0] # split weight matrix into blocks and randomly drop selected blocks mask = paddle.zeros( [in_features // block_size * out_features], dtype=paddle.bool) # the implementation of bernoulli_, p=0.5 mask = paddle.ones_like(mask) * 0.5 mask = paddle.bernoulli(mask) mask = mask.unsqueeze(1).tile([1, block_size]).reshape( [-1, in_features]) else: # gather weight and sizes weight = mod.weight in_channels = mod.weight.shape[1] out_channels = mod.weight.shape[0] # split weight matrix into blocks and randomly drop selected blocks if module.weight.shape[2:] == (1, 1): mask = paddle.zeros( [in_channels // block_size * out_channels], dtype=paddle.bool) # the implementation of bernoulli_, p=0.5 mask = paddle.ones_like(mask) * 0.5 mask = paddle.bernoulli(mask) mask = mask.unsqueeze(1).tile([1, block_size]).reshape( [-1, in_channels]) else: mask = paddle.zeros(weight.shape) # the implementation of bernoulli_, p=0.5 mask = paddle.ones_like(mask) * 0.5 mask = paddle.bernoulli(mask) mask = mask.unsqueeze(1).tile([1, in_channels, 1, 1]) # scale weights and apply mask s = 1 / (1 - p) mod.weight.set_value(s * weight.masked_fill(mask, 0)) module.register_forward_pre_hook(_forward_pre_hook) return module @with_incremental_state class MultiheadAttention(nn.Layer): """Multi-headed attention. See "Attention Is All You Need" for more details. """ def __init__( self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, self_attention=False, encoder_decoder_attention=False, q_noise=0.0, qn_block_size=8, # TODO: pass in config rather than string. # config defined in xformers.components.attention.AttentionConfig xformers_att_config: Optional[str]=None, xformers_blocksparse_layout: Optional[ paddle.Tensor]=None, # This should be part of the config xformers_blocksparse_blocksize: Optional[ int]=16, # This should be part of the config ): super().__init__() def eval_str_dict(x, type=dict): if x is None: return None if isinstance(x, str): x = eval(x) return x xformers_att_config = eval_str_dict(xformers_att_config) self.use_xformers = xformers_att_config is not None assert not self.use_xformers, "Do not use xformers in PaddleSpeech" self.embed_dim = embed_dim self.kdim = kdim if kdim is not None else embed_dim self.vdim = vdim if vdim is not None else embed_dim self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim self.num_heads = num_heads self.dropout_module = FairseqDropout( dropout, module_name=self.__class__.__name__) self.head_dim = embed_dim // num_heads assert (self.head_dim * num_heads == self.embed_dim ), "embed_dim must be divisible by num_heads" self.scaling = self.head_dim**-0.5 self.self_attention = self_attention self.encoder_decoder_attention = encoder_decoder_attention assert not self.self_attention or self.qkv_same_dim, ( "Self-attention requires query, key and " "value to be of the same size") # Todo scaled initialization # Empirically observed the convergence to be much better with # the scaled initialization weight_attr = nn.initializer.XavierUniform() kv_proj_bias_attr = nn.initializer.XavierUniform() out_proj_bias_attr = nn.initializer.Constant(0) self.k_proj = quant_noise( nn.Linear( self.kdim, embed_dim, weight_attr=weight_attr, bias_attr=bias if not bias else kv_proj_bias_attr), q_noise, qn_block_size) self.v_proj = quant_noise( nn.Linear( self.vdim, embed_dim, weight_attr=weight_attr, bias_attr=bias if not bias else kv_proj_bias_attr), q_noise, qn_block_size) self.q_proj = quant_noise( nn.Linear( embed_dim, embed_dim, weight_attr=weight_attr, bias_attr=bias), q_noise, qn_block_size) self.out_proj = quant_noise( nn.Linear( embed_dim, embed_dim, weight_attr=weight_attr, bias_attr=bias if not bias else out_proj_bias_attr), q_noise, qn_block_size) # nn.initializer.XavierUniform(self.k_proj.weight, gain=1 / math.sqrt(2)) # nn.initializer.XavierUniform(self.v_proj.weight, gain=1 / math.sqrt(2)) # nn.initializer.XavierUniform(self.q_proj.weight, gain=1 / math.sqrt(2)) # else: # self.k_proj.weight = paddle.ParamAttr() # nn.initializer.XavierUniform(self.k_proj.weight) # nn.initializer.XavierUniform(self.v_proj.weight) # nn.initializer.XavierUniform(self.q_proj.weight) # nn.initializer.XavierUniform(self.out_proj.weight) # if self.out_proj.bias is not None: # nn.initializer.Constant(self.out_proj.bias) # if self.bias_k is not None: # nn.initializer.XavierNormal(self.bias_k) # if self.bias_v is not None: # nn.initializer.XavierNormal(self.bias_v) # self.k_proj = Linear(self.kdim, embed_dim) # self.v_proj = Linear(self.vdim, embed_dim) # self.q_proj = Linear(embed_dim, embed_dim) # self.out_proj = Linear(embed_dim, embed_dim) if add_bias_kv: self.bias_k = paddle.create_parameter( shape=[1, 1, embed_dim], dtype='float32', initializer=nn.initializer.XavierUniform) self.bias_v = paddle.create_parameter( shape=[1, 1, embed_dim], dtype='float32', initializer=nn.initializer.XavierUniform) else: self.bias_k = self.bias_v = None self.add_zero_attn = add_zero_attn self.beam_size = 1 # self.reset_parameters() self.onnx_trace = False self.skip_embed_dim_check = False def prepare_for_onnx_export_(self): self.onnx_trace = True def reset_parameters(self): if self.qkv_same_dim: # Empirically observed the convergence to be much better with # the scaled initialization nn.initializer.XavierUniform( self.k_proj.weight, gain=1 / math.sqrt(2)) nn.initializer.XavierUniform( self.v_proj.weight, gain=1 / math.sqrt(2)) nn.initializer.XavierUniform( self.q_proj.weight, gain=1 / math.sqrt(2)) else: self.k_proj.weight = paddle.ParamAttr() nn.initializer.XavierUniform(self.k_proj.weight) nn.initializer.XavierUniform(self.v_proj.weight) nn.initializer.XavierUniform(self.q_proj.weight) nn.initializer.XavierUniform(self.out_proj.weight) if self.out_proj.bias is not None: nn.initializer.Constant(self.out_proj.bias) if self.bias_k is not None: nn.initializer.XavierNormal(self.bias_k) if self.bias_v is not None: nn.initializer.XavierNormal(self.bias_v) def _get_reserve_head_index(self, num_heads_to_keep: int): k_proj_heads_norm = [] q_proj_heads_norm = [] v_proj_heads_norm = [] for i in range(self.num_heads): start_idx = i * self.head_dim end_idx = (i + 1) * self.head_dim k_proj_heads_norm.append( paddle.sum( paddle.abs(self.k_proj.weight[:, start_idx:end_idx])) .tolist() + paddle.sum( paddle.abs(self.k_proj.bias[start_idx:end_idx])).tolist()) q_proj_heads_norm.append( paddle.sum( paddle.abs(self.q_proj.weight[:, start_idx:end_idx])) .tolist() + paddle.sum( paddle.abs(self.q_proj.bias[start_idx:end_idx])).tolist()) v_proj_heads_norm.append( paddle.sum( paddle.abs(self.v_proj.weight[:, start_idx:end_idx])) .tolist() + paddle.sum( paddle.abs(self.v_proj.bias[start_idx:end_idx])).tolist()) heads_norm = [] for i in range(self.num_heads): heads_norm.append(k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i]) sorted_head_index = sorted( range(self.num_heads), key=lambda k: heads_norm[k], reverse=True) reserve_head_index = [] for i in range(num_heads_to_keep): start = sorted_head_index[i] * self.head_dim end = (sorted_head_index[i] + 1) * self.head_dim reserve_head_index.append((start, end)) return reserve_head_index def _adaptive_prune_heads(self, reserve_head_index: List[Tuple[int, int]]): new_q_weight = [] new_q_bias = [] new_k_weight = [] new_k_bias = [] new_v_weight = [] new_v_bias = [] new_out_proj_weight = [] for ele in reserve_head_index: start_idx, end_idx = ele new_q_weight.append(self.q_proj.weight[:, start_idx:end_idx]) new_q_bias.append(self.q_proj.bias[start_idx:end_idx]) new_k_weight.append(self.k_proj.weight[:, start_idx:end_idx]) new_k_bias.append(self.k_proj.bias[start_idx:end_idx]) new_v_weight.append(self.v_proj.weight[:, start_idx:end_idx]) new_v_bias.append(self.v_proj.bias[start_idx:end_idx]) new_out_proj_weight.append( self.out_proj.weight[start_idx:end_idx, ]) new_q_weight = paddle.concat(new_q_weight, axis=-1).detach() new_k_weight = paddle.concat(new_k_weight, axis=-1).detach() new_v_weight = paddle.concat(new_v_weight, axis=-1).detach() new_out_proj_weight = paddle.concat(new_out_proj_weight).detach() new_q_weight.stop_gradient = False new_k_weight.stop_gradient = False new_v_weight.stop_gradient = False new_out_proj_weight.stop_gradient = False new_q_bias = paddle.concat(new_q_bias).detach() new_q_bias.stop_gradient = False new_k_bias = paddle.concat(new_k_bias).detach() new_k_bias.stop_gradient = False new_v_bias = paddle.concat(new_v_bias).detach() new_v_bias.stop_gradient = False self.q_proj.weight = paddle.create_parameter( shape=new_q_weight.shape, dtype=new_q_weight.dtype, default_initializer=paddle.nn.initializer.Assign(new_q_weight)) self.q_proj.bias = paddle.create_parameter( shape=new_q_bias.shape, dtype=new_q_bias.dtype, default_initializer=paddle.nn.initializer.Assign(new_q_bias)) self.k_proj.weight = paddle.create_parameter( shape=new_k_weight.shape, dtype=new_k_weight.dtype, default_initializer=paddle.nn.initializer.Assign(new_k_weight)) self.k_proj.bias = paddle.create_parameter( shape=new_k_bias.shape, dtype=new_k_bias.dtype, default_initializer=paddle.nn.initializer.Assign(new_k_bias)) self.v_proj.weight = paddle.create_parameter( shape=new_v_weight.shape, dtype=new_v_weight.dtype, default_initializer=paddle.nn.initializer.Assign(new_v_weight)) self.v_proj.bias = paddle.create_parameter( shape=new_v_bias.shape, dtype=new_v_bias.dtype, default_initializer=paddle.nn.initializer.Assign(new_v_bias)) self.out_proj.weight = paddle.create_parameter( shape=new_out_proj_weight.shape, dtype=new_out_proj_weight.dtype, default_initializer=paddle.nn.initializer.Assign( new_out_proj_weight)) self.num_heads = len(reserve_head_index) self.embed_dim = self.head_dim * self.num_heads self.q_proj.out_features = self.embed_dim self.k_proj.out_features = self.embed_dim self.v_proj.out_features = self.embed_dim def _set_skip_embed_dim_check(self): self.skip_embed_dim_check = True def _pad_masks( self, key_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor], ) -> Tuple[Optional[Tensor], Optional[Tensor]]: if attn_mask is not None: shape = attn_mask.shape[:-1] + [ 1, ] attn_mask = paddle.concat( [attn_mask, paddle.zeros(shape, dtype=attn_mask.dtype)], axis=-1) if key_padding_mask is not None: shape = key_padding_mask.shape[:-1] + [ 1, ] key_padding_mask = paddle.concat( [ key_padding_mask, paddle.zeros( shape, dtype=key_padding_mask.dtype) ], axis=-1) return key_padding_mask, attn_mask def _add_bias( self, k: Tensor, v: Tensor, key_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor], bsz: int, ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: assert self.bias_k is not None assert self.bias_v is not None k = paddle.concat([k, self.bias_k.tile([1, bsz, 1])], axis=-1) v = paddle.concat([v, self.bias_v.tile([1, bsz, 1])], axis=-1) key_padding_mask, attn_mask = self._pad_masks( key_padding_mask=key_padding_mask, attn_mask=attn_mask) return k, v, key_padding_mask, attn_mask def _append_zero_attn( self, k: Tensor, v: Tensor, key_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor], ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: zero_attn_shape = k.shape[:-2] + [1] + k.shape[-1:] k = paddle.concat( [k, paddle.zeros(zero_attn_shape, dtype=k.dtype)], axis=-2) v = paddle.concat( [v, paddle.zeros(zero_attn_shape, dtype=v.dtype)], axis=-2) key_padding_mask, attn_mask = self._pad_masks( key_padding_mask=key_padding_mask, attn_mask=attn_mask) return k, v, key_padding_mask, attn_mask def forward( self, query, key: Optional[Tensor], value: Optional[Tensor], key_padding_mask: Optional[Tensor]=None, incremental_state: Optional[Dict[str, Dict[str, Optional[ Tensor]]]]=None, need_weights: bool=True, static_kv: bool=False, attn_mask: Optional[Tensor]=None, before_softmax: bool=False, need_head_weights: bool=False, ) -> Tuple[Tensor, Optional[Tensor]]: """Input shape: Time x Batch x Channel Args: key_padding_mask (ByteTensor, optional): mask to exclude keys that are pads, of shape `(batch, src_len)`, where padding elements are indicated by 1s. need_weights (bool, optional): return the attention weights, averaged over heads (default: False). attn_mask (ByteTensor, optional): typically used to implement causal attention, where the mask prevents the attention from looking forward in time (default: None). before_softmax (bool, optional): return the raw attention weights and values before the attention softmax. need_head_weights (bool, optional): return the attention weights for each head. Implies *need_weights*. Default: return the average attention weights over all heads. """ if need_head_weights: need_weights = True is_tpu = query.place == "xla" tgt_len, bsz, embed_dim = query.shape src_len = tgt_len if not self.skip_embed_dim_check: assert (embed_dim == self.embed_dim ), f"query dim {embed_dim} != {self.embed_dim}" assert list(query.shape) == [tgt_len, bsz, embed_dim] if key is not None: src_len, key_bsz, _ = key.shape # if not torch.jit.is_scripting(): # assert value is not None # assert src_len, key_bsz == value.shape[:2] # if ( # not self.onnx_trace # and not is_tpu # don't use PyTorch version on TPUs # and incremental_state is None # and not static_kv # # A workaround for quantization to work. Otherwise JIT compilation # # treats bias in linear module as method. # and not torch.jit.is_scripting() # # The Multihead attention implemented in pytorch forces strong dimension check # # for input embedding dimention and K,Q,V projection dimension. # # Since pruning will break the dimension check and it is not easy to modify the pytorch API, # # it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check # and not self.skip_embed_dim_check # ): # assert key is not None and value is not None # if self.use_xformers: # return self._xformers_attn_forward( # query, key, value, key_padding_mask, need_weights, attn_mask # ) # else: # return F.multi_head_attention_forward( # query, # key, # value, # self.embed_dim, # self.num_heads, # torch.empty([0]), # torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)), # self.bias_k, # self.bias_v, # self.add_zero_attn, # self.dropout_module.p, # self.out_proj.weight, # self.out_proj.bias, # self.training or self.dropout_module.apply_during_inference, # key_padding_mask, # need_weights, # attn_mask, # use_separate_proj_weight=True, # q_proj_weight=self.q_proj.weight, # k_proj_weight=self.k_proj.weight, # v_proj_weight=self.v_proj.weight, # ) if incremental_state is not None: saved_state = self._get_input_buffer(incremental_state) if saved_state is not None and "prev_key" in saved_state: # previous time steps are cached - no need to recompute # key and value if they are static if static_kv: assert self.encoder_decoder_attention and not self.self_attention key = value = None else: saved_state = None if self.self_attention: q = self.q_proj(query) k = self.k_proj(query) v = self.v_proj(query) elif self.encoder_decoder_attention: # encoder-decoder attention q = self.q_proj(query) if key is None: assert value is None k = v = None else: if self.beam_size > 1 and bsz == key.size(1): # key is [T, bsz*beam_size, C], reduce to [T, bsz, C] key = key.view( key.size(0), -1, self.beam_size, key.size(2))[:, :, 0, :] if key_padding_mask is not None: key_padding_mask = key_padding_mask.view( -1, self.beam_size, key_padding_mask.size(1))[:, 0, :] k = self.k_proj(key) v = self.v_proj(key) else: assert key is not None and value is not None q = self.q_proj(query) k = self.k_proj(key) v = self.v_proj(value) q *= self.scaling if self.bias_k is not None: assert self.bias_v is not None k, v, attn_mask, key_padding_mask = self._add_bias( k, v, attn_mask, key_padding_mask, bsz) q = paddle.reshape( q, [tgt_len, bsz * self.num_heads, self.head_dim]).transpose( [1, 0, 2]) kv_bsz = bsz # need default value for scripting if k is not None: kv_bsz = k.shape[1] k = paddle.reshape( k, [-1, kv_bsz * self.num_heads, self.head_dim]).transpose( [1, 0, 2]) if v is not None: v = paddle.reshape( v, [-1, kv_bsz * self.num_heads, self.head_dim]).transpose( [1, 0, 2]) if saved_state is not None: # saved states are stored with shape (bsz, num_heads, seq_len, head_dim) if "prev_key" in saved_state: _prev_key = saved_state["prev_key"] assert _prev_key is not None kv_bsz = _prev_key.shape[0] prev_key = _prev_key.reshape( [kv_bsz * self.num_heads, -1, self.head_dim]) if static_kv: k = prev_key else: assert k is not None k = paddle.concat([prev_key, k], axis=1) src_len = k.shape[1] if "prev_value" in saved_state: _prev_value = saved_state["prev_value"] assert _prev_value is not None assert kv_bsz == _prev_value.size(0) prev_value = _prev_value.reshape( [kv_bsz * self.num_heads, -1, self.head_dim]) if static_kv: v = prev_value else: assert v is not None v = paddle.concat([prev_value, v], axis=1) prev_key_padding_mask: Optional[Tensor] = None if "prev_key_padding_mask" in saved_state: prev_key_padding_mask = saved_state["prev_key_padding_mask"] assert k is not None and v is not None key_padding_mask = MultiheadAttention._append_prev_key_padding_mask( key_padding_mask=key_padding_mask, prev_key_padding_mask=prev_key_padding_mask, batch_size=kv_bsz, src_len=k.shape[1], static_kv=static_kv, ) saved_state["prev_key"] = k.reshape( [kv_bsz, self.num_heads, -1, self.head_dim]) saved_state["prev_value"] = v.reshape( [kv_bsz, self.num_heads, -1, self.head_dim]) saved_state["prev_key_padding_mask"] = key_padding_mask # In this branch incremental_state is never None assert incremental_state is not None incremental_state = self._set_input_buffer(incremental_state, saved_state) assert k is not None assert k.shape[1] == src_len # This is part of a workaround to get around fork/join parallelism # not supporting Optional types. if key_padding_mask is not None and key_padding_mask.dim() == 0: key_padding_mask = None if key_padding_mask is not None: assert key_padding_mask.shape[0] == kv_bsz assert key_padding_mask.shape[1] == src_len if self.add_zero_attn: assert v is not None src_len += 1 k, v, key_padding_mask, attn_mask = self._append_zero_attn( k=k, v=v, key_padding_mask=key_padding_mask, attn_mask=attn_mask) if self.encoder_decoder_attention and bsz != kv_bsz: attn_weights = paddle.einsum( "bxhtd,bhsd->bxhts", q.reshape([kv_bsz, -1, self.num_heads] + q.shape[1:]), k.reshape([kv_bsz, self.num_heads] + k.shape[1:]), ) attn_weights = attn_weights.reshape([ -1, ] + attn_weights.shape[-2:]) else: attn_weights = paddle.bmm(q, k.transpose([0, 2, 1])) attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz) assert list( attn_weights.shape) == [bsz * self.num_heads, tgt_len, src_len] if attn_mask is not None: attn_mask = attn_mask.unsqueeze(0) if self.onnx_trace: attn_mask = attn_mask.tile([attn_weights.shape[0], 1, 1]) attn_weights += attn_mask if key_padding_mask is not None: # don't attend to padding symbols attn_weights = attn_weights.reshape( [bsz, self.num_heads, tgt_len, src_len]) if not is_tpu: attn_weights = attn_weights.reshape( [kv_bsz, -1, self.num_heads, tgt_len, src_len]) attn_weights = paddle.where( key_padding_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3) .astype('bool'), float('-inf') * paddle.ones_like(attn_weights), attn_weights) else: attn_weights = attn_weights.transpose([2, 1, 0]) attn_weights = paddle.where(key_padding_mask, float('-inf') * paddle.ones_like(attn_weights), attn_weights) attn_weights = attn_weights.transpose([2, 1, 0]) attn_weights = attn_weights.reshape( [bsz * self.num_heads, tgt_len, src_len]) if before_softmax: return attn_weights, v def softmax_supporting_onnx_trace(x, dim: int, onnx_trace: bool=False): if onnx_trace: return F.softmax(x, axis=dim) else: return F.softmax(x, axis=dim, dtype='float32') attn_weights_float = softmax_supporting_onnx_trace( attn_weights, dim=-1, onnx_trace=self.onnx_trace) attn_weights = paddle.cast(attn_weights_float, attn_weights.dtype) attn_probs = self.dropout_module(attn_weights) assert v is not None if self.encoder_decoder_attention and bsz != kv_bsz: attn = paddle.einsum( "bxhts,bhsd->bxhtd", attn_probs.reshape([kv_bsz, -1, self.num_heads] + attn_probs.shape[1:]), v.reshape([kv_bsz, self.num_heads] + v.shape[1:]), ) attn = attn.reshape([ -1, ] + attn.shape[-2:]) else: attn = paddle.bmm(attn_probs, v) assert list( attn.shape) == [bsz * self.num_heads, tgt_len, self.head_dim] if self.onnx_trace and attn.shape[1] == 1: # when ONNX tracing a single decoder step (sequence length == 1) # the transpose is a no-op copy before view, thus unnecessary attn = attn.reshape([tgt_len, bsz, self.embed_dim]) else: attn = attn.transpose([1, 0, 2]).reshape( [tgt_len, bsz, self.embed_dim]) attn = self.out_proj(attn) attn_weights: Optional[Tensor] = None if need_weights: attn_weights = attn_weights_float.reshape( [bsz, self.num_heads, tgt_len, src_len]).transpose([1, 0, 2, 3]) if not need_head_weights: # average attention weights over heads attn_weights = attn_weights.mean(axis=0) return attn, attn_weights @staticmethod def _append_prev_key_padding_mask( key_padding_mask: Optional[Tensor], prev_key_padding_mask: Optional[Tensor], batch_size: int, src_len: int, static_kv: bool, ) -> Optional[Tensor]: # saved key padding masks have shape (bsz, seq_len) if prev_key_padding_mask is not None and static_kv: new_key_padding_mask = prev_key_padding_mask elif prev_key_padding_mask is not None and key_padding_mask is not None: new_key_padding_mask = paddle.concat( [ paddle.cast(prev_key_padding_mask, 'float32'), paddle.cast(key_padding_mask, 'float32') ], axis=1) # During incremental decoding, as the padding token enters and # leaves the frame, there will be a time when prev or current # is None elif prev_key_padding_mask is not None: if src_len > prev_key_padding_mask.shape[1]: filler = paddle.zeros( [batch_size, src_len - prev_key_padding_mask.shape[1]], ) new_key_padding_mask = paddle.concat( [ paddle.cast(prev_key_padding_mask, 'float32'), paddle.cast(filler, 'float32') ], axis=1) else: new_key_padding_mask = prev_key_padding_mask elif key_padding_mask is not None: if src_len > key_padding_mask.shape[1]: filler = paddle.zeros( [batch_size, src_len - key_padding_mask.shape[1]], ) new_key_padding_mask = paddle.concat( [ paddle.cast(filler, 'float32'), paddle.cast(key_padding_mask, 'float32') ], axis=1) else: new_key_padding_mask = paddle.cast(key_padding_mask, 'float32') else: new_key_padding_mask = prev_key_padding_mask return new_key_padding_mask @paddle.jit.to_static def reorder_incremental_state( self, incremental_state: Dict[str, Dict[str, Optional[Tensor]]], new_order: Tensor, ): """Reorder buffered internal state (for incremental generation).""" input_buffer = self._get_input_buffer(incremental_state) if input_buffer is not None: for k in input_buffer.keys(): input_buffer_k = input_buffer[k] if input_buffer_k is not None: if self.encoder_decoder_attention: if input_buffer_k.shape[ 0] * self.beam_size == new_order.shape[0]: return incremental_state elif self.beam_size > 1: input_buffer[k] = paddle.index_select( input_buffer_k, index=new_order.reshape( [-1, self.beam_size])[:, 0] // self.beam_size, axis=0, ) else: input_buffer[k] = paddle.index_select( input_buffer_k, index=new_order, axis=0) else: input_buffer[k] = paddle.index_select( input_buffer_k, index=new_order, axis=0) incremental_state = self._set_input_buffer(incremental_state, input_buffer) return incremental_state def set_beam_size(self, beam_size): """Used for effiecient beamable enc-dec attention""" self.beam_size = beam_size def _get_input_buffer( self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] ) -> Dict[str, Optional[Tensor]]: result = self.get_incremental_state(incremental_state, "attn_state") if result is not None: return result else: empty_result: Dict[str, Optional[Tensor]] = {} return empty_result def _set_input_buffer( self, incremental_state: Dict[str, Dict[str, Optional[Tensor]]], buffer: Dict[str, Optional[Tensor]], ): return self.set_incremental_state(incremental_state, "attn_state", buffer) def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int): return attn_weights def upgrade_state_dict_named(self, state_dict, name): prefix = name + "." if name != "" else "" items_to_add = {} keys_to_remove = [] for k in state_dict.keys(): if k.endswith(prefix + "in_proj_weight"): # in_proj_weight used to be q + k + v with same dimensions dim = int(state_dict[k].shape[0] / 3) items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim] items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim:2 * dim] items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim:] keys_to_remove.append(k) k_bias = prefix + "in_proj_bias" if k_bias in state_dict.keys(): dim = int(state_dict[k].shape[0] / 3) items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim] items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][ dim:2 * dim] items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim:] keys_to_remove.append(prefix + "in_proj_bias") for k in keys_to_remove: del state_dict[k] for key, value in items_to_add.items(): state_dict[key] = value class GumbelVectorQuantizer(nn.Layer): def __init__( self, dim, num_vars, temp, groups, combine_groups, vq_dim, time_first, activation=nn.GELU(), weight_proj_depth=1, weight_proj_factor=1, ): """Vector quantization using gumbel softmax Args: dim: input dimension (channels) num_vars: number of quantized vectors per group temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor) groups: number of groups for vector quantization combine_groups: whether to use the vectors for all groups vq_dim: dimensionality of the resulting quantized vector time_first: if true, expect input in BxTxC format, otherwise in BxCxT activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1 weight_proj_depth: number of layers (with activation in between) to project input before computing logits weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of projections by this factor """ super().__init__() self.groups = groups self.combine_groups = combine_groups self.input_dim = dim self.num_vars = num_vars self.time_first = time_first assert ( vq_dim % groups == 0 ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation" var_dim = vq_dim // groups num_groups = groups if not combine_groups else 1 self.vars = self.create_parameter( (1, num_groups * num_vars, var_dim), default_initializer=nn.initializer.Uniform()) if weight_proj_depth > 1: def block(input_dim, output_dim): return nn.Sequential(Linear(input_dim, output_dim), activation) inner_dim = self.input_dim * weight_proj_factor self.weight_proj = nn.Sequential( *[ block(self.input_dim if i == 0 else inner_dim, inner_dim) for i in range(weight_proj_depth - 1) ], Linear(inner_dim, groups * num_vars), ) else: self.weight_proj = Linear( self.input_dim, groups * num_vars, weight_attr=nn.initializer.Normal(mean=0, std=1), bias_attr=nn.initializer.Zero()) if isinstance(temp, str): import ast temp = ast.literal_eval(temp) assert len(temp) == 3, f"{temp}, {len(temp)}" self.max_temp, self.min_temp, self.temp_decay = temp self.curr_temp = self.max_temp self.codebook_indices = None def set_num_updates(self, num_updates): self.curr_temp = max(self.max_temp * self.temp_decay**num_updates, self.min_temp) def get_codebook_indices(self): if self.codebook_indices is None: from itertools import product p = [range(self.num_vars)] * self.groups inds = list(product(*p)) self.codebook_indices = paddle.to_tensor( inds, dtype='int64', place=self.vars.place).flatten() if not self.combine_groups: self.codebook_indices = self.codebook_indices.reshape( self.num_vars**self.groups, -1) for b in range(1, self.groups): self.codebook_indices[:, b] += self.num_vars * b self.codebook_indices = self.codebook_indices.flatten() return self.codebook_indices def codebook(self): indices = self.get_codebook_indices() return (self.vars.squeeze(0).index_select(0, indices) .reshape(self.num_vars**self.groups, -1)) def sample_from_codebook(self, b, n): indices = self.get_codebook_indices() indices = indices.reshape(-1, self.groups) cb_size = indices.shape[0] assert (n < cb_size ), f"sample size {n} is greater than size of codebook {cb_size}" sample_idx = paddle.randint(low=0, high=cb_size, shape=(b * n, )) indices = indices[sample_idx] z = self.vars.squeeze(0).index_select(0, indices.flatten()).reshape( b, n, -1) return z def to_codebook_index(self, indices): res = paddle.full(indices.shape[:-1], 0, dtype=indices.dtype) for i in range(self.groups): exponent = self.groups - i - 1 res += indices[..., i] * (self.num_vars**exponent) return res def forward_idx(self, x): res = self.forward(x, produce_targets=True) return res["x"], res["targets"] def forward(self, x, produce_targets=False): result = {"num_vars": self.num_vars * self.groups} if not self.time_first: x = x.transpose([0, 2, 1]) bsz, tsz, fsz = x.shape x = x.reshape([-1, fsz]) x = self.weight_proj(x) x = x.reshape([bsz * tsz * self.groups, -1]) _, k = x.max(-1) hard_x = paddle.zeros_like(x) hard_x.scatter_(-1, k.reshape([-1, 1]), 1.0) hard_x = hard_x.reshape([bsz * tsz, self.groups, -1]) hard_probs = paddle.mean(hard_x.astype('float32'), axis=0) result["code_perplexity"] = paddle.exp(-paddle.sum( hard_probs * paddle.log(hard_probs + 1e-7), axis=-1)).sum() avg_probs = F.softmax( x.reshape([bsz * tsz, self.groups, -1]).astype('float32'), axis=-1).mean(axis=0) result["prob_perplexity"] = paddle.exp(-paddle.sum( avg_probs * paddle.log(avg_probs + 1e-7), axis=-1)).sum() result["temp"] = self.curr_temp if self.training: x = F.gumbel_softmax( x.astype('float32'), temperature=self.curr_temp, hard=True).astype(x.dtype) else: x = hard_x x = x.reshape([bsz * tsz, -1]) vars = self.vars if self.combine_groups: vars = vars.tile([1, self.groups, 1]) if produce_targets: result["targets"] = (x.reshape([bsz * tsz * self.groups, -1]) .argmax(axis=-1) .reshape([bsz, tsz, self.groups]).detach()) x = x.unsqueeze(-1) * vars x = x.reshape([bsz * tsz, self.groups, self.num_vars, -1]) x = x.sum(axis=-2) x = x.reshape([bsz, tsz, -1]) if not self.time_first: x = x.transpose([0, 2, 1]) result["x"] = x return result class GradMultiply(paddle.autograd.PyLayer): @staticmethod def forward(ctx, x, scale): ctx.scale = scale res = x.numpy().copy() return paddle.to_tensor(res, dtype=x.dtype) @staticmethod def backward(ctx, grad): return grad * ctx.scale, None class SamePad(nn.Layer): def __init__(self, kernel_size, causal=False): super().__init__() if causal: self.remove = kernel_size - 1 else: self.remove = 1 if kernel_size % 2 == 0 else 0 def forward(self, x): if self.remove > 0: x = x[:, :, :-self.remove] return x class TransposeLast(nn.Layer): def __init__(self, deconstruct_idx=None): super().__init__() self.deconstruct_idx = deconstruct_idx def forward(self, x): if self.deconstruct_idx is not None: x = x[self.deconstruct_idx] trans_dim = paddle.arange(x.dim()) trans_dim[-1], trans_dim[-2] = trans_dim[-2], trans_dim[-1] return x.transpose(trans_dim) class Fp32LayerNorm(LayerNorm): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def forward(self, input): output = F.layer_norm( input.astype('float32'), self._normalized_shape, self.weight.astype('float32') if self.weight is not None else None, self.bias.astype('float32') if self.bias is not None else None, self._epsilon, ) return output.astype(input.dtype) # Todo: change this when paddle supports F.group_norm class Fp32GroupNorm(nn.Layer): def __init__(self, *args, **kwargs): super().__init__() self.group_norm = paddle.nn.GroupNorm(*args, **kwargs) fp32_weight = paddle.create_parameter( shape=self.group_norm.weight.shape, dtype='float32', default_initializer=paddle.nn.initializer.Assign( self.group_norm.weight)) fp32_bias = paddle.create_parameter( shape=self.group_norm.bias.shape, dtype='float32', default_initializer=paddle.nn.initializer.Assign( self.group_norm.bias)) self.group_norm.weight = fp32_weight self.group_norm.bias = fp32_bias def forward(self, input): output = self.group_norm(input.astype('float32')) return output.astype(input.dtype) class StrEnumMeta(EnumMeta): # this is workaround for submitit pickling leading to instance checks failing in hydra for StrEnum, see # https://github.com/facebookresearch/hydra/issues/1156 @classmethod def __instancecheck__(cls, other): return "enum" in str(type(other)) class StrEnum(Enum, metaclass=StrEnumMeta): def __str__(self): return self.value def __eq__(self, other: str): return self.value == other def __repr__(self): return self.value def __hash__(self): return hash(str(self)) def ChoiceEnum(choices: List[str]): """return the Enum class used to enforce list of choices""" return StrEnum("Choices", {k: k for k in choices}) def relu_squared(x: paddle.Tensor): return F.relu(x).pow(2) def get_activation_fn(activation: str) -> Callable: """Returns the activation function corresponding to `activation`""" def gelu_accurate(x): if not hasattr(gelu_accurate, "_a"): gelu_accurate._a = math.sqrt(2 / math.pi) return (0.5 * x * (1 + paddle.tanh(gelu_accurate._a * (x + 0.044715 * paddle.pow(x, 3))))) def gelu(x: paddle.Tensor) -> paddle.Tensor: return paddle.nn.functional.gelu(x.astype('float32')).astype(x.dtype) if activation == "relu": return F.relu elif activation == "relu_squared": return relu_squared elif activation == "gelu": return gelu elif activation == "gelu_fast": return gelu_accurate elif activation == "gelu_accurate": return gelu_accurate elif activation == "tanh": return paddle.tanh elif activation == "linear": return lambda x: x elif activation == "swish": return paddle.nn.Swish else: raise RuntimeError( "--activation-fn {} not supported".format(activation)) def get_available_activation_fns() -> List: return [ "relu", "gelu", "gelu_fast", # deprecated "gelu_accurate", "tanh", "linear", ] def compute_mask_indices( shape: Tuple[int, int], padding_mask: Optional[paddle.Tensor], mask_prob: float, mask_length: int, mask_type: str="static", mask_other: float=0.0, min_masks: int=0, no_overlap: bool=False, min_space: int=0, require_same_masks: bool=True, mask_dropout: float=0.0, ) -> np.ndarray: """ Computes random mask spans for a given shape Args: shape: the the shape for which to compute masks. should be of size 2 where first element is batch size and 2nd is timesteps padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by number of timesteps divided by length of mask span to mask approximately this percentage of all elements. however due to overlaps, the actual number will be smaller (unless no_overlap is True) mask_type: how to compute mask lengths static = fixed size uniform = sample from uniform distribution [mask_other, mask_length*2] normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element poisson = sample from possion distribution with lambda = mask length min_masks: minimum number of masked spans no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample mask_dropout: randomly dropout this percentage of masks in each example """ bsz, all_sz = shape mask = np.full((bsz, all_sz), False) all_num_mask = int( # add a random number for probabilistic rounding mask_prob * all_sz / float(mask_length) + np.random.rand()) all_num_mask = max(min_masks, all_num_mask) mask_idcs = [] for i in range(bsz): if padding_mask is not None: sz = all_sz - padding_mask[i].long().sum().item() num_mask = int( # add a random number for probabilistic rounding mask_prob * sz / float(mask_length) + np.random.rand()) num_mask = max(min_masks, num_mask) else: sz = all_sz num_mask = all_num_mask if mask_type == "static": lengths = np.full(num_mask, mask_length) elif mask_type == "uniform": lengths = np.random.randint( mask_other, mask_length * 2 + 1, size=num_mask) elif mask_type == "normal": lengths = np.random.normal(mask_length, mask_other, size=num_mask) lengths = [max(1, int(round(x))) for x in lengths] elif mask_type == "poisson": lengths = np.random.poisson(mask_length, size=num_mask) lengths = [int(round(x)) for x in lengths] else: raise Exception("unknown mask selection " + mask_type) if sum(lengths) == 0: lengths[0] = min(mask_length, sz - 1) if no_overlap: mask_idc = [] def arrange(s, e, length, keep_length): span_start = np.random.randint(s, e - length) mask_idc.extend(span_start + i for i in range(length)) new_parts = [] if span_start - s - min_space >= keep_length: new_parts.append((s, span_start - min_space + 1)) if e - span_start - length - min_space > keep_length: new_parts.append((span_start + length + min_space, e)) return new_parts parts = [(0, sz)] min_length = min(lengths) for length in sorted(lengths, reverse=True): lens = np.fromiter( (e - s if e - s >= length + min_space else 0 for s, e in parts), np.int, ) l_sum = np.sum(lens) if l_sum == 0: break probs = lens / np.sum(lens) c = np.random.choice(len(parts), p=probs) s, e = parts.pop(c) parts.extend(arrange(s, e, length, min_length)) mask_idc = np.asarray(mask_idc) else: min_len = min(lengths) if sz - min_len <= num_mask: min_len = sz - num_mask - 1 mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) mask_idc = np.asarray([ mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j]) ]) mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) min_len = min([len(m) for m in mask_idcs]) for i, mask_idc in enumerate(mask_idcs): if len(mask_idc) > min_len and require_same_masks: mask_idc = np.random.choice(mask_idc, min_len, replace=False) if mask_dropout > 0: num_holes = np.rint(len(mask_idc) * mask_dropout).astype(int) mask_idc = np.random.choice( mask_idc, len(mask_idc) - num_holes, replace=False) mask[i, mask_idc] = True return mask def index_put(tensor, indices, value): tensor[indices] = value return tensor # ToDo if faster? def buffered_arange(max): if not hasattr(buffered_arange, "buf"): buffered_arange.buf = paddle.empty([max], dtype='int64') if max > buffered_arange.buf.numel(): buffered_arange.buf = paddle.arange(max) return buffered_arange.buf[:max] def pad_to_multiple(x, multiple, dim=-1, value=0): # Inspired from https://github.com/lucidrains/local-attention/blob/master/local_attention/local_attention.py#L41 if x is None: return None, 0 tsz = x.shape[dim] m = tsz / multiple remainder = math.ceil(m) * multiple - tsz if m.is_integer(): return x, 0 pad_offset = (0, ) * (-1 - dim) * 2 return F.pad( x, pad=[*pad_offset, 0, remainder, *pad_offset], value=value, data_format='NLC'), remainder EXTRACTOR_MODE_CHOICES = ChoiceEnum(["default", "layer_norm"]) MASKING_DISTRIBUTION_CHOICES = ChoiceEnum( ["static", "uniform", "normal", "poisson"]) LAYER_TYPE_CHOICES = ChoiceEnum(["transformer"]) # ToDo: conformer @dataclass class Wav2Vec2Config: extractor_mode: EXTRACTOR_MODE_CHOICES = field( default="default", metadata={ "help": "mode for feature extractor. default has a single group norm with d " "groups in the first conv block, whereas layer_norm has layer norms in " "every block (meant to use with normalize=True)" }, ) encoder_layers: int = field( default=12, metadata={"help": "num encoder layers in the transformer"}) encoder_embed_dim: int = field( default=768, metadata={"help": "encoder embedding dimension"}) encoder_ffn_embed_dim: int = field( default=3072, metadata={"help": "encoder embedding dimension for FFN"}) encoder_attention_heads: int = field( default=12, metadata={"help": "num encoder attention heads"}) activation_fn: ChoiceEnum(get_available_activation_fns()) = field( default="gelu", metadata={"help": "activation function to use"}) layer_type: LAYER_TYPE_CHOICES = field( default="transformer", metadata={"help": "layer type in encoder"}) # dropouts dropout: float = field( default=0.1, metadata={"help": "dropout probability for the transformer"}) attention_dropout: float = field( default=0.1, metadata={"help": "dropout probability for attention weights"}) activation_dropout: float = field( default=0.0, metadata={"help": "dropout probability after activation in FFN"}) encoder_layerdrop: float = field( default=0.0, metadata={"help": "probability of dropping a tarnsformer layer"}) dropout_input: float = field( default=0.0, metadata={"help": "dropout to apply to the input (after feat extr)"}, ) dropout_features: float = field( default=0.0, metadata={"help": "dropout to apply to the features (after feat extr)"}, ) final_dim: int = field( default=0, metadata={ "help": "project final representations and targets to this many dimensions." "set to encoder_embed_dim is <= 0" }, ) layer_norm_first: bool = field( default=False, metadata={"help": "apply layernorm first in the transformer"}) conv_feature_layers: str = field( default="[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512,2,2)] + [(512,2,2)]", metadata={ "help": "string describing convolutional feature extraction layers in form of a python list that contains " "[(dim, kernel_size, stride), ...]" }, ) conv_bias: bool = field( default=False, metadata={"help": "include bias in conv encoder"}) logit_temp: float = field( default=0.1, metadata={"help": "temperature to divide logits by"}) quantize_targets: bool = field( default=False, metadata={"help": "use quantized targets"}) quantize_input: bool = field( default=False, metadata={"help": "use quantized inputs"}) same_quantizer: bool = field( default=False, metadata={"help": "use same quantizer for inputs and targets"}) target_glu: bool = field( default=False, metadata={"help": "adds projection + glu to targets"}) feature_grad_mult: float = field( default=1.0, metadata={"help": "multiply feature extractor var grads by this"}) quantizer_depth: int = field( default=1, metadata={"help": "number of quantizer layers"}, ) quantizer_factor: int = field( default=3, metadata={ "help": "dimensionality increase for inner quantizer layers (if depth > 1)" }, ) latent_vars: int = field( default=320, metadata={ "help": "number of latent variables V in each group of the codebook" }, ) latent_groups: int = field( default=2, metadata={ "help": "number of groups G of latent variables in the codebook" }, ) latent_dim: int = field( default=0, metadata={ "help": "if > 0, uses this dimensionality for latent variables. " "otherwise uses final_dim / latent_groups" }, ) # masking mask_length: int = field(default=10, metadata={"help": "mask length"}) mask_prob: float = field( default=0.65, metadata={"help": "probability of replacing a token with mask"}) mask_selection: MASKING_DISTRIBUTION_CHOICES = field( default="static", metadata={"help": "how to choose mask length"}) mask_other: float = field( default=0, metadata={ "help": "secondary mask argument (used for more complex distributions), " "see help in compute_mask_indices" }, ) no_mask_overlap: bool = field( default=False, metadata={"help": "whether to allow masks to overlap"}) mask_min_space: int = field( default=1, metadata={"help": "min space between spans (if no overlap is enabled)"}, ) require_same_masks: bool = field( default=True, metadata={ "help": "whether to number of masked timesteps must be the same across all " "examples in a batch" }, ) mask_dropout: float = field( default=0.0, metadata={"help": "percent of masks to unmask for each sample"}, ) # channel masking mask_channel_length: int = field( default=10, metadata={"help": "length of the mask for features (channels)"}) mask_channel_prob: float = field( default=0.0, metadata={"help": "probability of replacing a feature with 0"}) mask_channel_before: bool = False mask_channel_selection: MASKING_DISTRIBUTION_CHOICES = field( default="static", metadata={"help": "how to choose mask length for channel masking"}, ) mask_channel_other: float = field( default=0, metadata={ "help": "secondary mask argument (used for more complex distributions), " "see help in compute_mask_indicesh" }, ) no_mask_channel_overlap: bool = field( default=False, metadata={"help": "whether to allow channel masks to overlap"}) mask_channel_min_space: int = field( default=1, metadata={"help": "min space between spans (if no overlap is enabled)"}, ) # negative selection num_negatives: int = field( default=100, metadata={"help": "number of negative examples from the same sample"}, ) negatives_from_everywhere: bool = field( default=False, metadata={ "help": "sample negatives from everywhere, not just masked states" }, ) cross_sample_negatives: int = field( default=0, metadata={"help": "number of negative examples from the any sample"}) codebook_negatives: int = field( default=0, metadata={"help": "number of negative examples codebook"}) # positional embeddings conv_pos: int = field( default=128, metadata={ "help": "number of filters for convolutional positional embeddings" }, ) conv_pos_groups: int = field( default=16, metadata={ "help": "number of groups for convolutional positional embedding" }, ) pos_conv_depth: int = field( default=1, metadata={"help": "depth of positional encoder network"}, ) latent_temp: Tuple[float, float, float] = field( default=(2, 0.5, 0.999995), metadata={ "help": "temperature for latent variable sampling. " "can be tuple of 3 values (start, end, decay)" }, ) max_positions: int = field( default=100000, metadata={"help": "Max positions"}) checkpoint_activations: bool = field( default=False, metadata={ "help": "recompute activations and save memory for extra compute" }, ) # FP16 optimization required_seq_len_multiple: int = field( default=2, metadata={ "help": "pad the input to encoder such that the sequence length is divisible by multiple" }, ) crop_seq_to_multiple: int = field( default=1, metadata={ "help": "crop convolutional feature extractor output such that the sequence length is divisible by multiple" }, ) # Conformer depthwise_conv_kernel_size: int = field( default=31, metadata={ "help": "depthwise-conv-kernel-size for convolution in conformer layer" }, ) attn_type: str = field( default="", metadata={"help": "if espnet use ESPNET MHA"}, ) pos_enc_type: str = field( default="abs", metadata={"help": "Positional encoding type to use in conformer"}, ) fp16: bool = field( default=False, metadata={"help": "If fp16 is being used"}) class Wav2Vec2Model(nn.Layer): def __init__(self, cfg: Wav2Vec2Config): super().__init__() self.cfg = cfg feature_enc_layers = eval(cfg.conv_feature_layers) self.embed = feature_enc_layers[-1][0] self.feature_extractor = ConvFeatureExtractionModel( conv_layers=feature_enc_layers, dropout=0.0, mode=cfg.extractor_mode, conv_bias=cfg.conv_bias, ) self.post_extract_proj = (Linear(self.embed, cfg.encoder_embed_dim) if self.embed != cfg.encoder_embed_dim and not cfg.quantize_input else None) self.crop_seq_to_multiple = cfg.crop_seq_to_multiple self.mask_prob = cfg.mask_prob self.mask_selection = cfg.mask_selection self.mask_other = cfg.mask_other self.mask_length = cfg.mask_length self.no_mask_overlap = cfg.no_mask_overlap self.mask_min_space = cfg.mask_min_space self.mask_channel_prob = cfg.mask_channel_prob self.mask_channel_before = cfg.mask_channel_before self.mask_channel_selection = cfg.mask_channel_selection self.mask_channel_other = cfg.mask_channel_other self.mask_channel_length = cfg.mask_channel_length self.no_mask_channel_overlap = cfg.no_mask_channel_overlap self.mask_channel_min_space = cfg.mask_channel_min_space self.dropout_input = nn.Dropout(cfg.dropout_input) self.dropout_features = nn.Dropout(cfg.dropout_features) self.feature_grad_mult = cfg.feature_grad_mult self.quantizer = None self.input_quantizer = None self.n_negatives = cfg.num_negatives self.cross_sample_negatives = cfg.cross_sample_negatives self.codebook_negatives = cfg.codebook_negatives self.negatives_from_everywhere = cfg.negatives_from_everywhere self.logit_temp = cfg.logit_temp final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim if cfg.quantize_targets: vq_dim = cfg.latent_dim if cfg.latent_dim > 0 else final_dim self.quantizer = GumbelVectorQuantizer( dim=self.embed, num_vars=cfg.latent_vars, temp=cfg.latent_temp, groups=cfg.latent_groups, combine_groups=False, vq_dim=vq_dim, time_first=True, weight_proj_depth=cfg.quantizer_depth, weight_proj_factor=cfg.quantizer_factor, ) self.project_q = Linear(vq_dim, final_dim) else: self.project_q = Linear(self.embed, final_dim) if cfg.quantize_input: if cfg.same_quantizer and self.quantizer is not None: vq_dim = final_dim self.input_quantizer = self.quantizer else: vq_dim = cfg.latent_dim if cfg.latent_dim > 0 else cfg.encoder_embed_dim self.input_quantizer = GumbelVectorQuantizer( dim=self.embed, num_vars=cfg.latent_vars, temp=cfg.latent_temp, groups=cfg.latent_groups, combine_groups=False, vq_dim=vq_dim, time_first=True, weight_proj_depth=cfg.quantizer_depth, weight_proj_factor=cfg.quantizer_factor, ) self.project_inp = Linear(vq_dim, cfg.encoder_embed_dim) self.mask_emb = self.create_parameter( shape=[cfg.encoder_embed_dim], default_initializer=paddle.nn.initializer.Uniform(), dtype='float32', ) encoder_cls = TransformerEncoder self.encoder = encoder_cls(cfg) self.layer_norm = LayerNorm(self.embed) self.target_glu = None if cfg.target_glu: self.target_glu = nn.Sequential( Linear(final_dim, final_dim * 2), GLU()) self.final_proj = Linear(cfg.encoder_embed_dim, final_dim) def upgrade_state_dict_named(self, state_dict, name): super().upgrade_state_dict_named(state_dict, name) """Upgrade a (possibly old) state dict for new versions of fairseq.""" return state_dict @classmethod def build_model(cls, cfg: Wav2Vec2Config, task=None): """Build a new model instance.""" return cls(cfg) def apply_mask( self, x, padding_mask, mask_indices=None, mask_channel_indices=None, ): B, T, C = x.shape if self.mask_channel_prob > 0 and self.mask_channel_before: mask_channel_indices = compute_mask_indices( (B, C), None, self.mask_channel_prob, self.mask_channel_length, self.mask_channel_selection, self.mask_channel_other, no_overlap=self.no_mask_channel_overlap, min_space=self.mask_channel_min_space, ) mask_channel_indices = ( paddle.to_tensor(mask_channel_indices, plcae=x.plcae) .unsqueeze(1).expand([-1, T, -1])) x[mask_channel_indices] = 0 if self.mask_prob > 0: if mask_indices is None: mask_indices = compute_mask_indices( (B, T), padding_mask, self.mask_prob, self.mask_length, self.mask_selection, self.mask_other, min_masks=2, no_overlap=self.no_mask_overlap, min_space=self.mask_min_space, require_same_masks=self.cfg.require_same_masks, mask_dropout=self.cfg.mask_dropout, ) mask_indices = paddle.to_tensor(mask_indices, place=x.place) x = index_put(x, mask_indices, self.mask_emb) else: mask_indices = None if self.mask_channel_prob > 0 and not self.mask_channel_before: if mask_channel_indices is None: mask_channel_indices = compute_mask_indices( (B, C), None, self.mask_channel_prob, self.mask_channel_length, self.mask_channel_selection, self.mask_channel_other, no_overlap=self.no_mask_channel_overlap, min_space=self.mask_channel_min_space, ) mask_channel_indices = ( paddle.to_tensor(mask_channel_indices, place=x.place) .unsqueeze(1).expand([-1, T, -1])) x = index_put(x, mask_channel_indices, 0) return x, mask_indices def sample_negatives(self, y, num, padding_count=None): if self.n_negatives == 0 and self.cross_sample_negatives == 0: return paddle.empty([0], dtype=y.dtype) bsz, tsz, fsz = y.shape y = y.reshape([-1, fsz]) # BTC => (BxT)C # FIXME: what happens if padding_count is specified? cross_high = tsz * bsz high = tsz - (padding_count or 0) with paddle.no_grad(): assert high > 1, f"{bsz,tsz,fsz}" if self.n_negatives > 0: tszs = (buffered_arange(num).unsqueeze(-1) .expand([-1, self.n_negatives]).flatten()) neg_idxs = paddle.randint( low=0, high=high - 1, shape=[bsz, self.n_negatives * num]) neg_idxs[neg_idxs >= tszs] += 1 if self.cross_sample_negatives > 0: tszs = (buffered_arange(num).unsqueeze(-1) .expand([-1, self.cross_sample_negatives]).flatten()) cross_neg_idxs = paddle.randint( low=0, high=cross_high - 1, shape=[bsz, self.cross_sample_negatives * num], ) cross_neg_idxs[cross_neg_idxs >= tszs] += 1 if self.n_negatives > 0: neg_idxs = neg_idxs + (paddle.arange(bsz).unsqueeze(1) * high) else: neg_idxs = cross_neg_idxs if self.cross_sample_negatives > 0 and self.n_negatives > 0: neg_idxs = paddle.concat([neg_idxs, cross_neg_idxs], axis=1) negs = y[neg_idxs.reshape([-1])] negs = negs.reshape( [bsz, num, self.n_negatives + self.cross_sample_negatives, fsz]).transpose([2, 0, 1, 3]) # to NxBxTxC return negs, neg_idxs def compute_preds(self, x, y, negatives): neg_is_pos = (y == negatives).all(-1) y = y.unsqueeze(0) targets = paddle.concat([y, negatives], axis=0) logits = paddle.nn.functional.cosine_similarity(x, targets, axis=-1) logits = logits / self.logit_temp logits = logits.astype(x.dtype) return logits def _get_feat_extract_output_lengths(self, input_lengths: paddle.Tensor): """ Computes the output length of the convolutional layers """ def _conv_out_length(input_length, kernel_size, stride): return paddle.floor((input_length - kernel_size) / stride + 1) conv_cfg_list = eval(self.cfg.conv_feature_layers) for i in range(len(conv_cfg_list)): input_lengths = _conv_out_length(input_lengths, conv_cfg_list[i][1], conv_cfg_list[i][2]) return paddle.cast(input_lengths, 'int64') def forward( self, source, padding_mask=None, mask=True, features_only=False, layer=None, mask_indices=None, mask_channel_indices=None, padding_count=None, ): if self.feature_grad_mult > 0: features = self.feature_extractor(source) if self.feature_grad_mult != 1.0: features = GradMultiply.apply(features, self.feature_grad_mult) else: with paddle.no_grad(): features = self.feature_extractor(source) features_pen = features.pow(2).mean() features = features.transpose([0, 2, 1]) features = self.layer_norm(features) unmasked_features = features.clone() if padding_mask is not None and padding_mask.any(): input_lengths = (1 - paddle.cast(padding_mask, 'int64')).sum(-1) # apply conv formula to get real output_lengths output_lengths = self._get_feat_extract_output_lengths( input_lengths) padding_mask = paddle.zeros( features.shape[:2], dtype=features.dtype) # these two operations makes sure that all values # before the output lengths indices are attended to padding_mask[(paddle.arange(padding_mask.shape[0]), output_lengths - 1, )] = 1 padding_mask = paddle.cast( (1 - padding_mask.flip([-1]).cumsum(-1).flip([-1])), 'bool') else: padding_mask = None time_steps_to_drop = features.shape[1] % self.crop_seq_to_multiple if time_steps_to_drop != 0: features = features[:, :-time_steps_to_drop] unmasked_features = unmasked_features[:, :-time_steps_to_drop] if padding_mask is not None: padding_mask = padding_mask[:, :-time_steps_to_drop] if self.post_extract_proj is not None: features = self.post_extract_proj(features) features = self.dropout_input(features) unmasked_features = self.dropout_features(unmasked_features) num_vars = None code_ppl = None prob_ppl = None curr_temp = None if self.input_quantizer: q = self.input_quantizer(features, produce_targets=False) features = q["x"] num_vars = q["num_vars"] code_ppl = q["code_perplexity"] prob_ppl = q["prob_perplexity"] curr_temp = q["temp"] features = self.project_inp(features) if mask: x, mask_indices = self.apply_mask( features, padding_mask, mask_indices=mask_indices, mask_channel_indices=mask_channel_indices, ) if mask_indices is not None: y = unmasked_features[mask_indices].reshape([ unmasked_features.shape[0], -1, unmasked_features.shape[-1] ]) else: x = features y = unmasked_features mask_indices = None x, layer_results = self.encoder( x, padding_mask=padding_mask, layer=layer) if features_only: return { "x": x, "padding_mask": padding_mask, "features": unmasked_features, "layer_results": layer_results, } if self.quantizer: if self.negatives_from_everywhere: q = self.quantizer(unmasked_features, produce_targets=False) y = q["x"] num_vars = q["num_vars"] code_ppl = q["code_perplexity"] prob_ppl = q["prob_perplexity"] curr_temp = q["temp"] y = self.project_q(y) negs, _ = self.sample_negatives( y, mask_indices[0].sum(), padding_count=padding_count, ) y = y[mask_indices].reshape([y.shape[0], -1, y.shape[-1]]) else: q = self.quantizer(y, produce_targets=False) y = q["x"] num_vars = q["num_vars"] code_ppl = q["code_perplexity"] prob_ppl = q["prob_perplexity"] curr_temp = q["temp"] y = self.project_q(y) negs, _ = self.sample_negatives( y, y.shape[1], padding_count=padding_count, ) if self.codebook_negatives > 0: cb_negs = self.quantizer.sample_from_codebook( y.shape[0] * y.shape[1], self.codebook_negatives) cb_negs = cb_negs.reshape( [self.codebook_negatives, y.shape[0], y.shape[1], -1]) # order doesnt matter cb_negs = self.project_q(cb_negs) negs = paddle.concat([negs, cb_negs], axis=0) else: y = self.project_q(y) if self.negatives_from_everywhere: negs, _ = self.sample_negatives( unmasked_features, y.shape[1], padding_count=padding_count, ) negs = self.project_q(negs) else: negs, _ = self.sample_negatives( y, y.shape[1], padding_count=padding_count, ) x = x[mask_indices].reshape([x.shape[0], -1, x.shape[-1]]) if self.target_glu: y = self.target_glu(y) negs = self.target_glu(negs) x = self.final_proj(x) x = self.compute_preds(x, y, negs) result = { "x": x, "padding_mask": padding_mask, "features_pen": features_pen, } if prob_ppl is not None: result["prob_perplexity"] = prob_ppl result["code_perplexity"] = code_ppl result["num_vars"] = num_vars result["temp"] = curr_temp return result def quantize(self, x): assert self.quantizer is not None x = self.feature_extractor(x) x = x.transpose([0, 2, 1]) x = self.layer_norm(x) return self.quantizer.forward_idx(x) def extract_features(self, source, padding_mask, mask=False, layer=None): res = self.forward( source, padding_mask, mask=mask, features_only=True, layer=layer) return res def get_logits(self, net_output): logits = net_output["x"] logits = logits.transpose([2, 1, 0]) logits = logits.reshape([-1, logits.shape[-1]]) return logits def get_targets(self, sample, net_output, expand_steps=True): x = net_output["x"] return paddle.zeros(x.shape[1] * x.shape[2], dtype='int64') def get_extra_losses(self, net_output): pen = [] if "prob_perplexity" in net_output: pen.append((net_output["num_vars"] - net_output["prob_perplexity"]) / net_output["num_vars"]) if "features_pen" in net_output: pen.append(net_output["features_pen"]) return pen def remove_pretraining_modules(self, last_layer=None): self.quantizer = None self.project_q = None self.target_glu = None self.final_proj = None if last_layer is not None: self.encoder.layers = nn.LayerList( l for i, l in enumerate(self.encoder.layers) if i <= last_layer) class ConvFeatureExtractionModel(nn.Layer): def __init__( self, conv_layers: List[Tuple[int, int, int]], dropout: float=0.0, mode: str="default", conv_bias: bool=False, ): super().__init__() assert mode in {"default", "layer_norm"} def block( n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False, ): def make_conv(): conv = Conv1D( n_in, n_out, k, stride=stride, bias_attr=conv_bias if not conv_bias else paddle.ParamAttr()) # nn.initializer.KaimingNormal()(conv.weight) return conv assert (is_layer_norm and is_group_norm ) is False, "layer norm and group norm are exclusive" if is_layer_norm: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), nn.Sequential( TransposeLast(), Fp32LayerNorm(dim), TransposeLast(), ), nn.GELU(), ) elif is_group_norm: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), Fp32GroupNorm(dim, dim), nn.GELU(), ) else: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), nn.GELU()) in_d = 1 self.conv_layers = nn.LayerList() for i, cl in enumerate(conv_layers): assert len(cl) == 3, "invalid conv definition: " + str(cl) (dim, k, stride) = cl self.conv_layers.append( block( in_d, dim, k, stride, is_layer_norm=mode == "layer_norm", is_group_norm=mode == "default" and i == 0, conv_bias=conv_bias, )) in_d = dim def forward(self, x): # BxT -> BxCxT x = x.unsqueeze(1) for conv in self.conv_layers: x = conv(x) return x def make_conv_pos(e, k, g): dropout = 0 std = math.sqrt((4 * (1.0 - dropout)) / (k * e)) pos_conv = Conv1D( e, e, kernel_size=k, padding=k // 2, groups=g, weight_attr=nn.initializer.Normal(mean=0, std=std), bias_attr=nn.initializer.Constant(0)) pos_conv = nn.utils.weight_norm(pos_conv, name="weight", dim=2) pos_conv = nn.Sequential(pos_conv, SamePad(k), nn.GELU()) return pos_conv class TransformerEncoder(nn.Layer): def build_encoder_layer(self, args: Wav2Vec2Config): layer = TransformerSentenceEncoderLayer( embedding_dim=self.embedding_dim, ffn_embedding_dim=args.encoder_ffn_embed_dim, num_attention_heads=args.encoder_attention_heads, dropout=self.dropout, attention_dropout=args.attention_dropout, activation_dropout=args.activation_dropout, activation_fn=args.activation_fn, layer_norm_first=args.layer_norm_first, ) return layer def __init__(self, args: Wav2Vec2Config): super().__init__() self.dropout = args.dropout self.embedding_dim = args.encoder_embed_dim self.required_seq_len_multiple = args.required_seq_len_multiple pos_conv_depth = getattr(args, "pos_conv_depth", 1) if pos_conv_depth > 1: num_layers = args.pos_conv_depth k = max(3, args.conv_pos // num_layers) def make_conv_block(e, k, g, l): return nn.Sequential(*[ nn.Sequential( Conv1D( e, e, kernel_size=k, padding=k // 2, groups=g, ), SamePad(k), TransposeLast(), LayerNorm(e, elementwise_affine=False), TransposeLast(), nn.GELU(), ) for _ in range(l) ]) self.pos_conv = make_conv_block(self.embedding_dim, k, args.conv_pos_groups, num_layers) else: self.pos_conv = make_conv_pos( self.embedding_dim, args.conv_pos, args.conv_pos_groups, ) self.layers = nn.LayerList([ self.build_encoder_layer(args) for _ in range(args.encoder_layers) ]) self.layer_norm_first = args.layer_norm_first self.layer_norm = LayerNorm(self.embedding_dim) self.layerdrop = args.encoder_layerdrop def forward(self, x, padding_mask=None, layer=None): x, layer_results = self.extract_features(x, padding_mask, layer) if self.layer_norm_first and layer is None: x = self.layer_norm(x) return x, layer_results def extract_features( self, x, padding_mask=None, tgt_layer=None, min_layer=0, ): if padding_mask is not None: x = index_put(x, padding_mask, 0) x_conv = self.pos_conv(x.transpose([0, 2, 1])) x_conv = x_conv.transpose([0, 2, 1]) x = x + x_conv if not self.layer_norm_first: x = self.layer_norm(x) # pad to the sequence length dimension x, pad_length = pad_to_multiple( x, self.required_seq_len_multiple, dim=-2, value=0) if pad_length > 0 and padding_mask is None: padding_mask = paddle.zeros([x.shape[0], x.shape[1]], dtype='bool') padding_mask[:, -pad_length:] = True else: padding_mask, _ = pad_to_multiple( padding_mask, self.required_seq_len_multiple, dim=-1, value=True) x = F.dropout(x, p=self.dropout, training=self.training) # B x T x C -> T x B x C x = x.transpose([1, 0, 2]) layer_results = [] r = None for i, layer in enumerate(self.layers): dropout_probability = np.random.random() if self.layerdrop > 0 else 1 if not self.training or (dropout_probability > self.layerdrop): x, (z, lr) = layer( x, self_attn_padding_mask=padding_mask, need_weights=False) if i >= min_layer: layer_results.append((x, z, lr)) if i == tgt_layer: r = x break if r is not None: x = r # T x B x C -> B x T x C x = x.transpose([1, 0, 2]) # undo paddding if pad_length > 0: x = x[:, :-pad_length] def undo_pad(a, b, c): return (a[:-pad_length], b[:-pad_length] if b is not None else b, c[:-pad_length], ) layer_results = [undo_pad(*u) for u in layer_results] return x, layer_results def max_positions(self): """Maximum output length supported by the encoder.""" return self.args.max_positions def upgrade_state_dict_named(self, state_dict, name): """Upgrade a (possibly old) state dict for new versions of fairseq.""" return state_dict class TransformerSentenceEncoderLayer(nn.Layer): """ Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained models. """ def __init__( self, embedding_dim: float=768, ffn_embedding_dim: float=3072, num_attention_heads: int=8, dropout: float=0.1, attention_dropout: float=0.1, activation_dropout: float=0.1, activation_fn: str="relu", layer_norm_first: bool=False, ) -> None: super().__init__() # Initialize parameters self.embedding_dim = embedding_dim self.dropout = dropout self.activation_dropout = activation_dropout # Initialize blocks self.activation_fn = get_activation_fn(activation_fn) self.self_attn = MultiheadAttention( self.embedding_dim, num_attention_heads, dropout=attention_dropout, self_attention=True, ) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(self.activation_dropout) self.dropout3 = nn.Dropout(dropout) self.layer_norm_first = layer_norm_first # layer norm associated with the self attention layer self.self_attn_layer_norm = LayerNorm(self.embedding_dim) self.fc1 = Linear(self.embedding_dim, ffn_embedding_dim) self.fc2 = Linear(ffn_embedding_dim, self.embedding_dim) # layer norm associated with the position wise feed-forward NN self.final_layer_norm = LayerNorm(self.embedding_dim) def forward( self, x: paddle.Tensor, self_attn_mask: paddle.Tensor=None, self_attn_padding_mask: paddle.Tensor=None, need_weights: bool=False, att_args=None, ): """ LayerNorm is applied either before or after the self-attention/ffn modules similar to the original Transformer imlementation. """ residual = x if self.layer_norm_first: x = self.self_attn_layer_norm(x) x, attn = self.self_attn( query=x, key=x, value=x, key_padding_mask=self_attn_padding_mask, attn_mask=self_attn_mask, need_weights=False, ) x = self.dropout1(x) x = residual + x residual = x x = self.final_layer_norm(x) x = self.activation_fn(self.fc1(x)) x = self.dropout2(x) x = self.fc2(x) layer_result = x x = self.dropout3(x) x = residual + x else: x, attn = self.self_attn( query=x, key=x, value=x, key_padding_mask=self_attn_padding_mask, need_weights=False, ) x = self.dropout1(x) x = residual + x x = self.self_attn_layer_norm(x) residual = x x = self.activation_fn(self.fc1(x)) x = self.dropout2(x) x = self.fc2(x) layer_result = x x = self.dropout3(x) x = residual + x x = self.final_layer_norm(x) return x, (attn, layer_result) @dataclass class AudioPretrainingConfig: sample_rate: int = field( default=16_000, metadata={ "help": "target sample rate. audio files will be up/down sampled to this rate" }, ) normalize: bool = field( default=False, metadata={ "help": "if set, normalizes input to have 0 mean and unit variance" }, ) enable_padding: bool = field( default=False, metadata={"help": "pad shorter samples instead of cropping"}) max_sample_size: Optional[int] = field( default=None, metadata={"help": "max sample size to crop to for batching"}) min_sample_size: Optional[int] = field( default=None, metadata={"help": "min sample size to skip small examples"})