



import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
       from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
       from tensorflow.keras.optimizers import Adam



class       UpperCAmelCase__(							tf.keras.optimizers.schedules.LearningRateSchedule	):

   '''simple docstring'''




   def __init__(						self  :			List[Any]			,				lowerCAmelCase  :			float			,				lowerCAmelCase  :			Callable			,				lowerCAmelCase  :			int			,				lowerCAmelCase  :			float = 1.0			,				lowerCAmelCase  :			str = None			,				)	->				List[Any]:




         """simple docstring"""

         super().__init__()
         lowercase__     =  initial_learning_rate
         lowercase__     =  warmup_steps
         lowercase__     =  power
         lowercase__     =  decay_schedule_fn
         lowercase__     =  name


   def __call__(						self  :			Optional[int]			,				lowerCAmelCase  :			str)	->				Any:




         """simple docstring"""

         with tf.name_scope(self.name or 'WarmUp') as name:
               # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
               # learning rate will be `global_step/num_warmup_steps * init_lr`.
               lowercase__     =  tf.cast(lowerCAmelCase			,				tf.floataa)
               lowercase__     =  tf.cast(self.warmup_steps			,				tf.floataa)
               lowercase__     =  global_step_float / warmup_steps_float
               lowercase__     =  self.initial_learning_rate * tf.math.pow(lowerCAmelCase			,				self.power)
               return tf.cond(
                   global_step_float < warmup_steps_float			,				lambda: warmup_learning_rate			,				lambda: self.decay_schedule_fn(step - self.warmup_steps)			,				name=lowerCAmelCase			,				)



   def 			UpperCAmelCase		(						self  :			Union[str, Any])	->				Optional[int]:




         """simple docstring"""

         return {
             "initial_learning_rate": self.initial_learning_rate,
             "decay_schedule_fn": self.decay_schedule_fn,
             "warmup_steps": self.warmup_steps,
             "power": self.power,
             "name": self.name,
         }

def   _lowerCAmelCase   (						A__						,  A__						,  A__						,  A__ = 0.0						,  A__ = 0.9						,  A__ = 0.9_99						,  A__ = 1E-8						,  A__ = None						,  A__ = None						,  A__ = 0.0						,  A__ = 1.0						,  A__ = None						,  ):
      lowercase__     =  tf.keras.optimizers.schedules.PolynomialDecay(
          initial_learning_rate=A__						,  decay_steps=num_train_steps - num_warmup_steps						,  end_learning_rate=init_lr * min_lr_ratio						,  power=A__						,  )
      if num_warmup_steps:
            lowercase__     =  WarmUp(
                initial_learning_rate=A__						,  decay_schedule_fn=A__						,  warmup_steps=A__						,  )
      if weight_decay_rate > 0.0:
            lowercase__     =  AdamWeightDecay(
                learning_rate=A__						,  weight_decay_rate=A__						,  beta_a=A__						,  beta_a=A__						,  epsilon=A__						,  clipnorm=A__						,  global_clipnorm=A__						,  exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias']						,  include_in_weight_decay=A__						,  )
      else:
            lowercase__     =  tf.keras.optimizers.Adam(
                learning_rate=A__						,  beta_a=A__						,  beta_a=A__						,  epsilon=A__						,  clipnorm=A__						,  global_clipnorm=A__						,  )
      # We return the optimizer and the LR scheduler in order to better track the
      # evolution of the LR independently of the optimizer.
      return optimizer, lr_schedule



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def __init__(						self  :			Tuple			,				lowerCAmelCase  :			Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.0_01			,				lowerCAmelCase  :			float = 0.9			,				lowerCAmelCase  :			float = 0.9_99			,				lowerCAmelCase  :			float = 1E-7			,				lowerCAmelCase  :			bool = False			,				lowerCAmelCase  :			float = 0.0			,				lowerCAmelCase  :			Optional[List[str]] = None			,				lowerCAmelCase  :			Optional[List[str]] = None			,				lowerCAmelCase  :			str = "AdamWeightDecay"			,				**lowerCAmelCase  :			Dict			,				)	->				List[str]:




         """simple docstring"""

         super().__init__(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				**lowerCAmelCase)
         lowercase__     =  weight_decay_rate
         lowercase__     =  include_in_weight_decay
         lowercase__     =  exclude_from_weight_decay


   @classmethod
   def 			UpperCAmelCase		(						cls  :			Union[str, Any]			,				lowerCAmelCase  :			Optional[int])	->				Tuple:




         """simple docstring"""

         lowercase__     =  {'WarmUp': WarmUp}
         return super(lowerCAmelCase			,				cls).from_config(lowerCAmelCase			,				custom_objects=lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			str			,				lowerCAmelCase  :			Tuple			,				lowerCAmelCase  :			Dict			,				lowerCAmelCase  :			Any)	->				Any:




         """simple docstring"""

         super(lowerCAmelCase			,				self)._prepare_local(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase)
         lowercase__     =  tf.constant(
             self.weight_decay_rate			,				name='adam_weight_decay_rate')


   def 			UpperCAmelCase		(						self  :			List[str]			,				lowerCAmelCase  :			Union[str, Any]			,				lowerCAmelCase  :			Tuple			,				lowerCAmelCase  :			Optional[Any])	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  self._do_use_weight_decay(var.name)
         if do_decay:
               return var.assign_sub(
                   learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['weight_decay_rate']			,				use_locking=self._use_locking			,				)
         return tf.no_op()


   def 			UpperCAmelCase		(						self  :			Tuple			,				lowerCAmelCase  :			Optional[int]			,				lowerCAmelCase  :			Optional[int]=None			,				**lowerCAmelCase  :			Optional[Any])	->				Dict:




         """simple docstring"""

         lowercase__,	lowercase__     =  list(zip(*lowerCAmelCase))
         return super(lowerCAmelCase			,				self).apply_gradients(zip(lowerCAmelCase			,				lowerCAmelCase)			,				name=lowerCAmelCase			,				**lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Union[str, Any]			,				lowerCAmelCase  :			Union[str, Any]			,				lowerCAmelCase  :			Dict			,				lowerCAmelCase  :			Tuple)	->				Tuple:




         """simple docstring"""

         if apply_state is None:
               return self._decayed_lr_t[var_dtype], {}

         lowercase__     =  apply_state or {}
         lowercase__     =  apply_state.get((var_device, var_dtype))
         if coefficients is None:
               lowercase__     =  self._fallback_apply_state(lowerCAmelCase			,				lowerCAmelCase)
               lowercase__     =  coefficients

         return coefficients["lr_t"], {"apply_state": apply_state}


   def 			UpperCAmelCase		(						self  :			Optional[Any]			,				lowerCAmelCase  :			List[str]			,				lowerCAmelCase  :			Dict			,				lowerCAmelCase  :			Union[str, Any]=None)	->				str:




         """simple docstring"""

         lowercase__,	lowercase__     =  self._get_lr(var.device			,				var.dtype.base_dtype			,				lowerCAmelCase)
         lowercase__     =  self._decay_weights_op(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase)
         with tf.control_dependencies([decay]):
               return super(lowerCAmelCase			,				self)._resource_apply_dense(lowerCAmelCase			,				lowerCAmelCase			,				**lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Tuple			,				lowerCAmelCase  :			int			,				lowerCAmelCase  :			List[Any]			,				lowerCAmelCase  :			Any			,				lowerCAmelCase  :			Optional[int]=None)	->				List[Any]:




         """simple docstring"""

         lowercase__,	lowercase__     =  self._get_lr(var.device			,				var.dtype.base_dtype			,				lowerCAmelCase)
         lowercase__     =  self._decay_weights_op(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase)
         with tf.control_dependencies([decay]):
               return super(lowerCAmelCase			,				self)._resource_apply_sparse(lowerCAmelCase			,				lowerCAmelCase			,				lowerCAmelCase			,				**lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			List[str])	->				Union[str, Any]:




         """simple docstring"""

         lowercase__     =  super().get_config()
         config.update({'weight_decay_rate': self.weight_decay_rate})
         return config



   def 			UpperCAmelCase		(						self  :			Union[str, Any]			,				lowerCAmelCase  :			Optional[Any])	->				Optional[int]:




         """simple docstring"""

         if self.weight_decay_rate == 0:
               return False

         if self._include_in_weight_decay:
               for r in self._include_in_weight_decay:
                     if re.search(lowerCAmelCase			,				lowerCAmelCase) is not None:
                           return True

         if self._exclude_from_weight_decay:
               for r in self._exclude_from_weight_decay:
                     if re.search(lowerCAmelCase			,				lowerCAmelCase) is not None:
                           return False
         return True



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def __init__(						self  :			List[Any])	->				List[Any]:




         """simple docstring"""

         lowercase__     =  []
         lowercase__     =  None


   @property
   def 			UpperCAmelCase		(						self  :			Any)	->				Tuple:




         """simple docstring"""

         if self._accum_steps is None:
               lowercase__     =  tf.Variable(
                   tf.constant(0			,				dtype=tf.intaa)			,				trainable=lowerCAmelCase			,				synchronization=tf.VariableSynchronization.ON_READ			,				aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA			,				)

         return self._accum_steps.value()


   @property
   def 			UpperCAmelCase		(						self  :			List[Any])	->				Dict:




         """simple docstring"""

         if not self._gradients:
               raise ValueError('The accumulator should be called first to initialize the gradients')
         return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]


   def __call__(						self  :			str			,				lowerCAmelCase  :			Any)	->				Optional[Any]:




         """simple docstring"""

         if not self._gradients:
               lowercase__     =  self.step  # Create the step variable.
               self._gradients.extend(
                   [
                       tf.Variable(
                           tf.zeros_like(lowerCAmelCase)			,				trainable=lowerCAmelCase			,				synchronization=tf.VariableSynchronization.ON_READ			,				aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA			,				)
                       if gradient is not None
                       else gradient
                       for gradient in gradients
                   ])
         if len(lowerCAmelCase) != len(self._gradients):
               raise ValueError(f'''Expected {len(self._gradients)} gradients, but got {len(lowerCAmelCase)}''')

         for accum_gradient, gradient in zip(self._gradients			,				lowerCAmelCase):
               if accum_gradient is not None and gradient is not None:
                     accum_gradient.assign_add(lowerCAmelCase)

         self._accum_steps.assign_add(1)



   def 			UpperCAmelCase		(						self  :			List[str])	->				List[Any]:




         """simple docstring"""

         if not self._gradients:
               return
         self._accum_steps.assign(0)
         for gradient in self._gradients:
               if gradient is not None:
                     gradient.assign(tf.zeros_like(lowerCAmelCase))





from ...configuration_utils import PretrainedConfig
from ...utils import logging


a__      :				Dict   	=			logging.get_logger(__name__)

a__      :				List[Any]   	=			{
    "facebook/s2t-small-librispeech-asr": (
        "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/config.json"
    ),
    # See all Speech2Text models at https://huggingface.co/models?filter=speech_to_text
}



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   A    :  int       				=    "speech_to_text"
   A    :  Optional[Any]       				=    ["past_key_values"]
   A    :  Optional[int]       				=    {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}


   def __init__(						self  :			Optional[int]			,				lowerCAmelCase  :			Tuple=1_00_00			,				lowerCAmelCase  :			int=12			,				lowerCAmelCase  :			int=20_48			,				lowerCAmelCase  :			Union[str, Any]=4			,				lowerCAmelCase  :			str=6			,				lowerCAmelCase  :			Dict=20_48			,				lowerCAmelCase  :			Dict=4			,				lowerCAmelCase  :			Optional[int]=0.0			,				lowerCAmelCase  :			Union[str, Any]=0.0			,				lowerCAmelCase  :			int=True			,				lowerCAmelCase  :			Optional[Any]=True			,				lowerCAmelCase  :			Dict="relu"			,				lowerCAmelCase  :			Tuple=2_56			,				lowerCAmelCase  :			Dict=0.1			,				lowerCAmelCase  :			Optional[Any]=0.0			,				lowerCAmelCase  :			List[Any]=0.0			,				lowerCAmelCase  :			Any=0.02			,				lowerCAmelCase  :			List[Any]=2			,				lowerCAmelCase  :			Tuple=True			,				lowerCAmelCase  :			Tuple=1			,				lowerCAmelCase  :			List[str]=0			,				lowerCAmelCase  :			Union[str, Any]=2			,				lowerCAmelCase  :			Any=60_00			,				lowerCAmelCase  :			Optional[int]=10_24			,				lowerCAmelCase  :			Optional[Any]=2			,				lowerCAmelCase  :			Optional[Any]=(5, 5)			,				lowerCAmelCase  :			Union[str, Any]=10_24			,				lowerCAmelCase  :			List[Any]=80			,				lowerCAmelCase  :			List[str]=1			,				**lowerCAmelCase  :			List[str]			,				)	->				Dict:




         """simple docstring"""

         lowercase__     =  vocab_size
         lowercase__     =  d_model
         lowercase__     =  encoder_ffn_dim
         lowercase__     =  encoder_layers
         lowercase__     =  encoder_attention_heads
         lowercase__     =  decoder_ffn_dim
         lowercase__     =  decoder_layers
         lowercase__     =  decoder_attention_heads
         lowercase__     =  dropout
         lowercase__     =  attention_dropout
         lowercase__     =  activation_dropout
         lowercase__     =  activation_function
         lowercase__     =  init_std
         lowercase__     =  encoder_layerdrop
         lowercase__     =  decoder_layerdrop
         lowercase__     =  use_cache
         lowercase__     =  encoder_layers
         lowercase__     =  scale_embedding  # scale factor will be sqrt(d_model) if True
         lowercase__     =  max_source_positions
         lowercase__     =  max_target_positions
         lowercase__     =  num_conv_layers
         lowercase__     =  list(lowerCAmelCase)
         lowercase__     =  conv_channels
         lowercase__     =  input_feat_per_channel
         lowercase__     =  input_channels

         if len(self.conv_kernel_sizes) != self.num_conv_layers:
               raise ValueError(
                   'Configuration for convolutional module is incorrect. '
                   'It is required that `len(config.conv_kernel_sizes)` == `config.num_conv_layers` '
                   f'''but is `len(config.conv_kernel_sizes) = {len(self.conv_kernel_sizes)}`, '''
                   f'''`config.num_conv_layers = {self.num_conv_layers}`.''')

         super().__init__(
             pad_token_id=lowerCAmelCase			,				bos_token_id=lowerCAmelCase			,				eos_token_id=lowerCAmelCase			,				is_encoder_decoder=lowerCAmelCase			,				decoder_start_token_id=lowerCAmelCase			,				**lowerCAmelCase			,				)
