"""simple docstring"""







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  _lowerCamelCase			(tf.keras.optimizers.schedules.LearningRateSchedule     ):







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

					"""simple docstring"""






					super().__init__()
					_lowercase :    Dict							=							initial_learning_rate
					_lowercase :    List[Any]							=							warmup_steps
					_lowercase :    Tuple							=							power
					_lowercase :    int							=							decay_schedule_fn
					_lowercase :    Any							=							name







				def __call__( self    :       List[str]			,					lowerCamelCase_    :       int						):

					"""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 :    Optional[Any]							=							tf.cast(lowerCamelCase_			,					tf.floataa						)
						_lowercase :    Optional[Any]							=							tf.cast(self.warmup_steps			,					tf.floataa						)
						_lowercase :    List[str]							=							global_step_float / warmup_steps_float
						_lowercase :    Dict							=							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    :       Dict						):

					"""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(		__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase = 0.0     ,__UpperCAmelCase = 0.9     ,__UpperCAmelCase = 0.9_9_9     ,__UpperCAmelCase = 1E-8     ,__UpperCAmelCase = None     ,__UpperCAmelCase = None     ,__UpperCAmelCase = 0.0     ,__UpperCAmelCase = 1.0     ,__UpperCAmelCase = None     ,):






	"""simple docstring"""





	_lowercase :    str							=							tf.keras.optimizers.schedules.PolynomialDecay(
	    initial_learning_rate=__UpperCAmelCase     ,decay_steps=num_train_steps - num_warmup_steps     ,end_learning_rate=init_lr * min_lr_ratio     ,power=__UpperCAmelCase     ,)
	if num_warmup_steps:
		_lowercase :    List[str]							=							WarmUp(
		    initial_learning_rate=__UpperCAmelCase     ,decay_schedule_fn=__UpperCAmelCase     ,warmup_steps=__UpperCAmelCase     ,)
	if weight_decay_rate > 0.0:
		_lowercase :    Union[str, Any]							=							AdamWeightDecay(
		    learning_rate=__UpperCAmelCase     ,weight_decay_rate=__UpperCAmelCase     ,beta_a=__UpperCAmelCase     ,beta_a=__UpperCAmelCase     ,epsilon=__UpperCAmelCase     ,clipnorm=__UpperCAmelCase     ,global_clipnorm=__UpperCAmelCase     ,exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias']     ,include_in_weight_decay=__UpperCAmelCase     ,)
	else:
		_lowercase :    Optional[int]							=							tf.keras.optimizers.Adam(
		    learning_rate=__UpperCAmelCase     ,beta_a=__UpperCAmelCase     ,beta_a=__UpperCAmelCase     ,epsilon=__UpperCAmelCase     ,clipnorm=__UpperCAmelCase     ,global_clipnorm=__UpperCAmelCase     ,)
	# 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  _lowerCamelCase			(__lowerCamelCase     ):







				def __init__( self    :       Dict			,					lowerCamelCase_    :       Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001			,					lowerCamelCase_    :       float = 0.9			,					lowerCamelCase_    :       float = 0.999			,					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_    :       int			,					):

					"""simple docstring"""






					super().__init__(lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_			,					**lowerCamelCase_						)
					_lowercase :    Union[str, Any]							=							weight_decay_rate
					_lowercase :    str							=							include_in_weight_decay
					_lowercase :    Dict							=							exclude_from_weight_decay







				@classmethod
				def 							__UpperCAmelCase					( cls    :       Union[str, Any]			,					lowerCamelCase_    :       str						):

					"""simple docstring"""






					_lowercase :    Tuple							=							{'WarmUp': WarmUp}
					return super(lowerCamelCase_			,					cls						).from_config(lowerCamelCase_			,					custom_objects=lowerCamelCase_						)







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       List[Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Dict						):

					"""simple docstring"""






					super(lowerCamelCase_			,					self						)._prepare_local(lowerCamelCase_			,					lowerCamelCase_			,					lowerCamelCase_						)
					_lowercase :    Union[str, Any]							=							tf.constant(
					    self.weight_decay_rate			,					name='adam_weight_decay_rate'						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Tuple						):

					"""simple docstring"""






					_lowercase :    Any							=							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    :       Optional[int]			,					lowerCamelCase_    :       Any			,					lowerCamelCase_    :       Any=None			,					**lowerCamelCase_    :       Dict						):

					"""simple docstring"""






					_lowercase :    int							=							list(zip(*lowerCamelCase_						)						)
					return super(lowerCamelCase_			,					self						).apply_gradients(zip(lowerCamelCase_			,					lowerCamelCase_						)			,					name=lowerCamelCase_			,					**lowerCamelCase_						)







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       str						):

					"""simple docstring"""






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

					_lowercase :    int							=							apply_state or {}
					_lowercase :    List[Any]							=							apply_state.get((var_device, var_dtype)						)
					if coefficients is None:
						_lowercase :    Any							=							self._fallback_apply_state(lowerCamelCase_			,					lowerCamelCase_						)
						_lowercase :    Union[str, Any]							=							coefficients

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







				def 							__UpperCAmelCase					( self    :       Union[str, Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Optional[Any]			,					lowerCamelCase_    :       Optional[int]=None						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							self._get_lr(var.device			,					var.dtype.base_dtype			,					lowerCamelCase_						)
					_lowercase :    Union[str, Any]							=							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    :       Any			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       int			,					lowerCamelCase_    :       Union[str, Any]			,					lowerCamelCase_    :       Union[str, Any]=None						):

					"""simple docstring"""






					_lowercase :    Any							=							self._get_lr(var.device			,					var.dtype.base_dtype			,					lowerCamelCase_						)
					_lowercase :    List[Any]							=							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    :       Dict						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							super().get_config()
					config.update({'weight_decay_rate': self.weight_decay_rate}						)
					return config







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Optional[Any]						):

					"""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  _lowerCamelCase			(__lowerCamelCase     ):







				def __init__( self    :       Union[str, Any]						):

					"""simple docstring"""






					_lowercase :    Any							=							[]
					_lowercase :    List[Any]							=							None







				@property
				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					if self._accum_steps is None:
						_lowercase :    List[str]							=							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    :       Optional[Any]						):

					"""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    :       List[Any]			,					lowerCamelCase_    :       int						):

					"""simple docstring"""






					if not self._gradients:
						_lowercase :    Optional[int]							=							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[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_						)						)




"""simple docstring"""







from typing import List, Optional, Union

import numpy as np

from ....audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
from ....feature_extraction_sequence_utils import SequenceFeatureExtractor
from ....feature_extraction_utils import BatchFeature
from ....file_utils import PaddingStrategy, TensorType
from ....utils import logging


SCREAMING_SNAKE_CASE									=						logging.get_logger(__name__)



class  _lowerCamelCase			(__lowerCamelCase     ):
				_snake_case     =  ["input_features", "attention_mask"]







				def __init__( self    :       int			,					lowerCamelCase_    :       List[str]=8_0			,					lowerCamelCase_    :       Tuple=1_6_0_0_0			,					lowerCamelCase_    :       Union[str, Any]=0.0			,					lowerCamelCase_    :       List[Any]=1_0			,					lowerCamelCase_    :       List[str]=2_5			,					lowerCamelCase_    :       List[Any]="hamming_window"			,					lowerCamelCase_    :       Tuple=3_2768.0			,					lowerCamelCase_    :       int=0.97			,					lowerCamelCase_    :       Optional[int]=1.0			,					lowerCamelCase_    :       Optional[Any]=True			,					lowerCamelCase_    :       Union[str, Any]=True			,					lowerCamelCase_    :       Optional[int]=False			,					**lowerCamelCase_    :       Optional[Any]			,					):

					"""simple docstring"""






					super().__init__(feature_size=lowerCamelCase_			,					sampling_rate=lowerCamelCase_			,					padding_value=lowerCamelCase_			,					**lowerCamelCase_						)

					_lowercase :    Dict							=							feature_size
					_lowercase :    Dict							=							sampling_rate
					_lowercase :    Tuple							=							padding_value
					_lowercase :    int							=							hop_length
					_lowercase :    Any							=							win_length
					_lowercase :    Union[str, Any]							=							frame_signal_scale
					_lowercase :    Tuple							=							preemphasis_coeff
					_lowercase :    Tuple							=							mel_floor
					_lowercase :    Tuple							=							normalize_means
					_lowercase :    List[Any]							=							normalize_vars
					_lowercase :    List[str]							=							win_function
					_lowercase :    int							=							return_attention_mask

					_lowercase :    Optional[Any]							=							win_length * sampling_rate // 1_0_0_0
					_lowercase :    Tuple							=							hop_length * sampling_rate // 1_0_0_0

					_lowercase :    str							=							optimal_fft_length(self.sample_size						)
					_lowercase :    Dict							=							(self.n_fft // 2) + 1







				def 							__UpperCAmelCase					( self    :       List[str]			,					lowerCamelCase_    :       np.array						):

					"""simple docstring"""






					if self.win_function == "hamming_window":
						_lowercase :    List[Any]							=							window_function(window_length=self.sample_size			,					name=self.win_function			,					periodic=lowerCamelCase_						)
					else:
						_lowercase :    Union[str, Any]							=							window_function(window_length=self.sample_size			,					name=self.win_function						)

					_lowercase :    Tuple							=							mel_filter_bank(
					    num_frequency_bins=self.n_freqs			,					num_mel_filters=self.feature_size			,					min_frequency=0.0			,					max_frequency=self.sampling_rate / 2.0			,					sampling_rate=self.sampling_rate			,					)

					_lowercase :    Tuple							=							spectrogram(
					    one_waveform * self.frame_signal_scale			,					window=lowerCamelCase_			,					frame_length=self.sample_size			,					hop_length=self.sample_stride			,					fft_length=self.n_fft			,					center=lowerCamelCase_			,					preemphasis=self.preemphasis_coeff			,					mel_filters=lowerCamelCase_			,					mel_floor=self.mel_floor			,					log_mel='log'			,					)
					return msfc_features.T







				def 							__UpperCAmelCase					( self    :       Optional[Any]			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Dict			,					lowerCamelCase_    :       Tuple						):

					"""simple docstring"""






					if self.normalize_means:
						_lowercase :    Optional[int]							=							x[:input_length].mean(axis=0						)
						_lowercase :    int							=							np.subtract(lowerCamelCase_			,					lowerCamelCase_						)
					if self.normalize_vars:
						_lowercase :    int							=							x[:input_length].std(axis=0						)
						_lowercase :    Optional[Any]							=							np.divide(lowerCamelCase_			,					lowerCamelCase_						)

					if input_length < x.shape[0]:
						_lowercase :    Dict							=							padding_value

					# make sure array is in float32
					_lowercase :    Tuple							=							x.astype(np.floataa						)

					return x







				def 							__UpperCAmelCase					( self    :       List[Any]			,					lowerCamelCase_    :       List[np.ndarray]			,					lowerCamelCase_    :       Optional[np.ndarray] = None						):

					"""simple docstring"""






					_lowercase :    Dict							=							attention_mask.sum(-1						) if attention_mask is not None else [x.shape[0] for x in input_features]
					return [self._normalize_one(lowerCamelCase_			,					lowerCamelCase_			,					self.padding_value						) for x, n in zip(lowerCamelCase_			,					lowerCamelCase_						)]







				def __call__( self    :       Dict			,					lowerCamelCase_    :       Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]			,					lowerCamelCase_    :       Union[bool, str, PaddingStrategy] = False			,					lowerCamelCase_    :       Optional[int] = None			,					lowerCamelCase_    :       bool = False			,					lowerCamelCase_    :       Optional[int] = None			,					lowerCamelCase_    :       Optional[bool] = None			,					lowerCamelCase_    :       Optional[Union[str, TensorType]] = None			,					lowerCamelCase_    :       Optional[int] = None			,					**lowerCamelCase_    :       Optional[int]			,					):

					"""simple docstring"""






					if sampling_rate is not None:
						if sampling_rate != self.sampling_rate:
							raise ValueError(
							    F'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
							    F''' {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with'''
							    F''' {self.sampling_rate} and not {sampling_rate}.'''						)
					else:
						logger.warning(
						    'It is strongly recommended to pass the ``sampling_rate`` argument to this function. '
						    'Failing to do so can result in silent errors that might be hard to debug.'						)

					_lowercase :    Optional[Any]							=							isinstance(lowerCamelCase_			,					np.ndarray						) and len(raw_speech.shape						) > 1
					if is_batched_numpy and len(raw_speech.shape						) > 2:
						raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''						)
					_lowercase :    Optional[int]							=							is_batched_numpy or (
					    isinstance(lowerCamelCase_			,					(list, tuple)						) and (isinstance(raw_speech[0]			,					(np.ndarray, tuple, list)						))
					)

					if is_batched:
						_lowercase :    str							=							[np.asarray(lowerCamelCase_			,					dtype=np.floataa						) for speech in raw_speech]
					elif not is_batched and not isinstance(lowerCamelCase_			,					np.ndarray						):
						_lowercase :    Tuple							=							np.asarray(lowerCamelCase_			,					dtype=np.floataa						)
					elif isinstance(lowerCamelCase_			,					np.ndarray						) and raw_speech.dtype is np.dtype(np.floataa						):
						_lowercase :    Optional[int]							=							raw_speech.astype(np.floataa						)

					# always return batch
					if not is_batched:
						_lowercase :    str							=							[raw_speech]

					# extract fbank features
					_lowercase :    Optional[Any]							=							[self._extract_mfsc_features(lowerCamelCase_						) for one_waveform in raw_speech]

					# convert into correct format for padding
					_lowercase :    Optional[int]							=							BatchFeature({'input_features': features}						)

					_lowercase :    Tuple							=							self.pad(
					    lowerCamelCase_			,					padding=lowerCamelCase_			,					max_length=lowerCamelCase_			,					truncation=lowerCamelCase_			,					pad_to_multiple_of=lowerCamelCase_			,					return_attention_mask=lowerCamelCase_			,					**lowerCamelCase_			,					)
					# make sure list is in array format
					_lowercase :    Dict							=							padded_inputs.get('input_features'						)
					if isinstance(input_features[0]			,					lowerCamelCase_						):
						_lowercase :    List[str]							=							[np.asarray(lowerCamelCase_			,					dtype=np.floataa						) for feature in input_features]

					_lowercase :    List[Any]							=							padded_inputs.get('attention_mask'						)
					if attention_mask is not None:
						_lowercase :    Union[str, Any]							=							[np.asarray(lowerCamelCase_			,					dtype=np.intaa						) for array in attention_mask]

					if self.normalize_means or self.normalize_vars:
						_lowercase :    int							=							(
						    np.array(lowerCamelCase_			,					dtype=np.intaa						)
						    if self._get_padding_strategies(lowerCamelCase_			,					max_length=lowerCamelCase_						) is not PaddingStrategy.DO_NOT_PAD
						    and padding
						    else None
						)
						_lowercase :    List[Any]							=							self.normalize(
						    padded_inputs['input_features']			,					attention_mask=lowerCamelCase_						)

					if return_tensors is not None:
						_lowercase :    Union[str, Any]							=							padded_inputs.convert_to_tensors(lowerCamelCase_						)

					return padded_inputs



