


"""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				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase = 1.0    ,				__UpperCAmelCase = None    ,				)       ->							List[Any]:
          super().__init__()
          _a    	=				initial_learning_rate
          _a    	=				warmup_steps
          _a    	=				power
          _a    	=				decay_schedule_fn
          _a    	=				name


     def __call__(      self    ,				__UpperCAmelCase							)       ->							Optional[int]:
          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`.
               _a    	=				tf.cast(__UpperCAmelCase    ,				tf.floataa							)
               _a    	=				tf.cast(self.warmup_steps    ,				tf.floataa							)
               _a    	=				global_step_float / warmup_steps_float
               _a    	=				self.initial_learning_rate * tf.math.pow(__UpperCAmelCase    ,				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=__UpperCAmelCase    ,				)





     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          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       A_   (   _lowerCAmelCase     :					float,     _lowerCAmelCase     :					int,     _lowerCAmelCase     :					int,     _lowerCAmelCase     :					float = 0.0,     _lowerCAmelCase     :					float = 0.9,     _lowerCAmelCase     :					float = 0.9_9_9,     _lowerCAmelCase     :					float = 1e-8,     _lowerCAmelCase     :					Optional[float] = None,     _lowerCAmelCase     :					Optional[float] = None,     _lowerCAmelCase     :					float = 0.0,     _lowerCAmelCase     :					float = 1.0,     _lowerCAmelCase     :					Optional[List[str]] = None,     ):






     """simple docstring"""

     _a    	=				tf.keras.optimizers.schedules.PolynomialDecay(
         initial_learning_rate=_lowerCAmelCase,     decay_steps=num_train_steps - num_warmup_steps,     end_learning_rate=init_lr * min_lr_ratio,     power=_lowerCAmelCase,     )
     if num_warmup_steps:
          _a    	=				WarmUp(
              initial_learning_rate=_lowerCAmelCase,     decay_schedule_fn=_lowerCAmelCase,     warmup_steps=_lowerCAmelCase,     )
     if weight_decay_rate > 0.0:
          _a    	=				AdamWeightDecay(
              learning_rate=_lowerCAmelCase,     weight_decay_rate=_lowerCAmelCase,     beta_a=_lowerCAmelCase,     beta_a=_lowerCAmelCase,     epsilon=_lowerCAmelCase,     clipnorm=_lowerCAmelCase,     global_clipnorm=_lowerCAmelCase,     exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias'''],     include_in_weight_decay=_lowerCAmelCase,     )
     else:
          _a    	=				tf.keras.optimizers.Adam(
              learning_rate=_lowerCAmelCase,     beta_a=_lowerCAmelCase,     beta_a=_lowerCAmelCase,     epsilon=_lowerCAmelCase,     clipnorm=_lowerCAmelCase,     global_clipnorm=_lowerCAmelCase,     )
     # 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  (					a__				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase = 0.001    ,				__UpperCAmelCase = 0.9    ,				__UpperCAmelCase = 0.999    ,				__UpperCAmelCase = 1e-7    ,				__UpperCAmelCase = False    ,				__UpperCAmelCase = 0.0    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = "AdamWeightDecay"    ,				**__UpperCAmelCase    ,				)       ->							Any:
          super().__init__(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				**__UpperCAmelCase							)
          _a    	=				weight_decay_rate
          _a    	=				include_in_weight_decay
          _a    	=				exclude_from_weight_decay


     @classmethod
     def 							_UpperCAmelCase      (      cls    ,				__UpperCAmelCase							)       ->							Optional[int]:
          _a    	=				{'''WarmUp''': WarmUp}
          return super(__UpperCAmelCase    ,				cls							).from_config(__UpperCAmelCase    ,				custom_objects=__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Union[str, Any]:
          super(__UpperCAmelCase    ,				self							)._prepare_local(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
          _a    	=				tf.constant(
              self.weight_decay_rate    ,				name='''adam_weight_decay_rate'''							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Union[str, Any]:
          _a    	=				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    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None    ,				**__UpperCAmelCase							)       ->							Any:
          _a	,      _a    	=				list(zip(*__UpperCAmelCase							)							)
          return super(__UpperCAmelCase    ,				self							).apply_gradients(zip(__UpperCAmelCase    ,				__UpperCAmelCase							)    ,				name=__UpperCAmelCase    ,				**__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							int:
          if apply_state is None:
               return self._decayed_lr_t[var_dtype], {}

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

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


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None							)       ->							Optional[int]:
          _a	,      _a    	=				self._get_lr(var.device    ,				var.dtype.base_dtype    ,				__UpperCAmelCase							)
          _a    	=				self._decay_weights_op(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
          with tf.control_dependencies([decay]							):
               return super(__UpperCAmelCase    ,				self							)._resource_apply_dense(__UpperCAmelCase    ,				__UpperCAmelCase    ,				**__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None							)       ->							str:
          _a	,      _a    	=				self._get_lr(var.device    ,				var.dtype.base_dtype    ,				__UpperCAmelCase							)
          _a    	=				self._decay_weights_op(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
          with tf.control_dependencies([decay]							):
               return super(__UpperCAmelCase    ,				self							)._resource_apply_sparse(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				**__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self							)       ->							str:
          _a    	=				super().get_config()
          config.update({'''weight_decay_rate''': self.weight_decay_rate}							)
          return config





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[Any]:
          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(__UpperCAmelCase    ,				__UpperCAmelCase							) is not None:
                         return True

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







class     __lowerCamelCase  (					a__				):

     '''simple docstring'''


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


     @property
     def 							_UpperCAmelCase      (      self							)       ->							Dict:
          if self._accum_steps is None:
               _a    	=				tf.Variable(
                   tf.constant(0    ,				dtype=tf.intaa							)    ,				trainable=__UpperCAmelCase    ,				synchronization=tf.VariableSynchronization.ON_READ    ,				aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA    ,				)

          return self._accum_steps.value()


     @property
     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          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    ,				__UpperCAmelCase							)       ->							List[str]:
          if not self._gradients:
               _a    	=				self.step  # Create the step variable.
               self._gradients.extend(
                   [
                       tf.Variable(
                           tf.zeros_like(__UpperCAmelCase							)    ,				trainable=__UpperCAmelCase    ,				synchronization=tf.VariableSynchronization.ON_READ    ,				aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA    ,				)
                       if gradient is not None
                       else gradient
                       for gradient in gradients
                   ]							)
          if len(__UpperCAmelCase							) != len(self._gradients							):
               raise ValueError(F'Expected {len(self._gradients							)} gradients, but got {len(__UpperCAmelCase							)}'							)

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

          self._accum_steps.assign_add(1							)





     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          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(__UpperCAmelCase							)							)



"""simple docstring"""






from .glue import glue_convert_examples_to_features, glue_output_modes, glue_processors, glue_tasks_num_labels
from .squad import SquadExample, SquadFeatures, SquadVaProcessor, SquadVaProcessor, squad_convert_examples_to_features
from .utils import DataProcessor, InputExample, InputFeatures, SingleSentenceClassificationProcessor
from .xnli import xnli_output_modes, xnli_processors, xnli_tasks_num_labels