


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				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE = 1.0				,       _SCREAMING_SNAKE_CASE = None				,       )->Tuple:




         '''simple docstring'''





         super().__init__()
         A_      :				List[str]		     =					initial_learning_rate
         A_      :				Optional[int]		     =					warmup_steps
         A_      :				Any		     =					power
         A_      :				Union[str, Any]		     =					decay_schedule_fn
         A_      :				Tuple		     =					name



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







       def        _snake_case       (      self		)->List[str]:




         '''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   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE = 0.0 ,			SCREAMING_SNAKE_CASE = 0.9 ,			SCREAMING_SNAKE_CASE = 0.9_9_9 ,			SCREAMING_SNAKE_CASE = 1e-8 ,			SCREAMING_SNAKE_CASE = None ,			SCREAMING_SNAKE_CASE = None ,			SCREAMING_SNAKE_CASE = 0.0 ,			SCREAMING_SNAKE_CASE = 1.0 ,			SCREAMING_SNAKE_CASE = None ,			):
  A_      :				Dict		     =					tf.keras.optimizers.schedules.PolynomialDecay(
      initial_learning_rate=SCREAMING_SNAKE_CASE ,			decay_steps=num_train_steps - num_warmup_steps ,			end_learning_rate=init_lr * min_lr_ratio ,			power=SCREAMING_SNAKE_CASE ,			)
  if num_warmup_steps:
    A_      :				Tuple		     =					WarmUp(
        initial_learning_rate=SCREAMING_SNAKE_CASE ,			decay_schedule_fn=SCREAMING_SNAKE_CASE ,			warmup_steps=SCREAMING_SNAKE_CASE ,			)
  if weight_decay_rate > 0.0:
    A_      :				Optional[int]		     =					AdamWeightDecay(
        learning_rate=SCREAMING_SNAKE_CASE ,			weight_decay_rate=SCREAMING_SNAKE_CASE ,			beta_a=SCREAMING_SNAKE_CASE ,			beta_a=SCREAMING_SNAKE_CASE ,			epsilon=SCREAMING_SNAKE_CASE ,			clipnorm=SCREAMING_SNAKE_CASE ,			global_clipnorm=SCREAMING_SNAKE_CASE ,			exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias'''] ,			include_in_weight_decay=SCREAMING_SNAKE_CASE ,			)
  else:
    A_      :				Optional[int]		     =					tf.keras.optimizers.Adam(
        learning_rate=SCREAMING_SNAKE_CASE ,			beta_a=SCREAMING_SNAKE_CASE ,			beta_a=SCREAMING_SNAKE_CASE ,			epsilon=SCREAMING_SNAKE_CASE ,			clipnorm=SCREAMING_SNAKE_CASE ,			global_clipnorm=SCREAMING_SNAKE_CASE ,			)
  # 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 (					UpperCamelCase			):




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE = 0.0_0_1				,       _SCREAMING_SNAKE_CASE = 0.9				,       _SCREAMING_SNAKE_CASE = 0.9_9_9				,       _SCREAMING_SNAKE_CASE = 1e-7				,       _SCREAMING_SNAKE_CASE = False				,       _SCREAMING_SNAKE_CASE = 0.0				,       _SCREAMING_SNAKE_CASE = None				,       _SCREAMING_SNAKE_CASE = None				,       _SCREAMING_SNAKE_CASE = "AdamWeightDecay"				,       **_SCREAMING_SNAKE_CASE				,       )->Optional[Any]:




         '''simple docstring'''





         super().__init__(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)
         A_      :				List[str]		     =					weight_decay_rate
         A_      :				Tuple		     =					include_in_weight_decay
         A_      :				Optional[int]		     =					exclude_from_weight_decay



       @classmethod
       def        _snake_case       (      cls				,       _SCREAMING_SNAKE_CASE		)->Optional[int]:




         '''simple docstring'''





         A_      :				str		     =					{'''WarmUp''': WarmUp}
         return super(_SCREAMING_SNAKE_CASE				,       cls		).from_config(_SCREAMING_SNAKE_CASE				,       custom_objects=_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->Optional[Any]:




         '''simple docstring'''





         super(_SCREAMING_SNAKE_CASE				,       self		)._prepare_local(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Optional[Any]		     =					tf.constant(
             self.weight_decay_rate				,       name='''adam_weight_decay_rate'''		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->Tuple:




         '''simple docstring'''





         A_      :				Optional[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        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=None				,       **_SCREAMING_SNAKE_CASE		)->Tuple:




         '''simple docstring'''





         A_     ,       A_      :				List[str]		     =					list(zip(*_SCREAMING_SNAKE_CASE		)		)
         return super(_SCREAMING_SNAKE_CASE				,       self		).apply_gradients(zip(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)				,       name=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->Optional[Any]:




         '''simple docstring'''





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

         A_      :				Tuple		     =					apply_state or {}
         A_      :				Optional[int]		     =					apply_state.get((var_device, var_dtype)		)
         if coefficients is None:
           A_      :				List[Any]		     =					self._fallback_apply_state(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
           A_      :				List[Any]		     =					coefficients

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



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=None		)->Any:




         '''simple docstring'''





         A_     ,       A_      :				Tuple		     =					self._get_lr(var.device				,       var.dtype.base_dtype				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Optional[Any]		     =					self._decay_weights_op(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         with tf.control_dependencies([decay]		):
           return super(_SCREAMING_SNAKE_CASE				,       self		)._resource_apply_dense(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=None		)->Dict:




         '''simple docstring'''





         A_     ,       A_      :				int		     =					self._get_lr(var.device				,       var.dtype.base_dtype				,       _SCREAMING_SNAKE_CASE		)
         A_      :				Dict		     =					self._decay_weights_op(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         with tf.control_dependencies([decay]		):
           return super(_SCREAMING_SNAKE_CASE				,       self		)._resource_apply_sparse(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





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







       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE		)->List[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(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		) is not None:
               return True

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






class   _lowerCamelCase (					UpperCamelCase			):




       """simple docstring"""



       def __init__(      self		)->Any:




         '''simple docstring'''





         A_      :				Tuple		     =					[]
         A_      :				Union[str, Any]		     =					None



       @property
       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         if self._accum_steps is None:
           A_      :				Optional[int]		     =					tf.Variable(
               tf.constant(0				,       dtype=tf.intaa		)				,       trainable=_SCREAMING_SNAKE_CASE				,       synchronization=tf.VariableSynchronization.ON_READ				,       aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA				,       )

         return self._accum_steps.value()



       @property
       def        _snake_case       (      self		)->Tuple:




         '''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				,       _SCREAMING_SNAKE_CASE		)->Tuple:




         '''simple docstring'''





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

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

         self._accum_steps.assign_add(1		)







       def        _snake_case       (      self		)->List[str]:




         '''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(_SCREAMING_SNAKE_CASE		)		)








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase								=	{
    """configuration_roberta""": ["""ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """RobertaConfig""", """RobertaOnnxConfig"""],
    """tokenization_roberta""": ["""RobertaTokenizer"""],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=	["""RobertaTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=	[
         """ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """RobertaForCausalLM""",
         """RobertaForMaskedLM""",
         """RobertaForMultipleChoice""",
         """RobertaForQuestionAnswering""",
         """RobertaForSequenceClassification""",
         """RobertaForTokenClassification""",
         """RobertaModel""",
         """RobertaPreTrainedModel""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=	[
         """TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """TFRobertaForCausalLM""",
         """TFRobertaForMaskedLM""",
         """TFRobertaForMultipleChoice""",
         """TFRobertaForQuestionAnswering""",
         """TFRobertaForSequenceClassification""",
         """TFRobertaForTokenClassification""",
         """TFRobertaMainLayer""",
         """TFRobertaModel""",
         """TFRobertaPreTrainedModel""",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase								=	[
         """FlaxRobertaForCausalLM""",
         """FlaxRobertaForMaskedLM""",
         """FlaxRobertaForMultipleChoice""",
         """FlaxRobertaForQuestionAnswering""",
         """FlaxRobertaForSequenceClassification""",
         """FlaxRobertaForTokenClassification""",
         """FlaxRobertaModel""",
         """FlaxRobertaPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaOnnxConfig
     from .tokenization_roberta import RobertaTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_roberta_fast import RobertaTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_roberta import (
              ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
              RobertaForCausalLM,
              RobertaForMaskedLM,
              RobertaForMultipleChoice,
              RobertaForQuestionAnswering,
              RobertaForSequenceClassification,
              RobertaForTokenClassification,
              RobertaModel,
              RobertaPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_roberta import (
              TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFRobertaForCausalLM,
              TFRobertaForMaskedLM,
              TFRobertaForMultipleChoice,
              TFRobertaForQuestionAnswering,
              TFRobertaForSequenceClassification,
              TFRobertaForTokenClassification,
              TFRobertaMainLayer,
              TFRobertaModel,
              TFRobertaPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_roberta import (
              FlaxRobertaForCausalLM,
              FlaxRobertaForMaskedLM,
              FlaxRobertaForMultipleChoice,
              FlaxRobertaForQuestionAnswering,
              FlaxRobertaForSequenceClassification,
              FlaxRobertaForTokenClassification,
              FlaxRobertaModel,
              FlaxRobertaPreTrainedModel,
          )

else:
     import sys

     UpperCamelCase								=	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




