





"""simple docstring"""


import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  from transformers.models.roberta.modeling_flax_roberta import (
      FlaxRobertaForCausalLM,
      FlaxRobertaForMaskedLM,
      FlaxRobertaForMultipleChoice,
      FlaxRobertaForQuestionAnswering,
      FlaxRobertaForSequenceClassification,
      FlaxRobertaForTokenClassification,
      FlaxRobertaModel,
  )



class   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


  def __init__(		self :							Tuple       , __a :							Any       , __a :							Optional[Any]=13       , __a :							Dict=7       , __a :							List[str]=True       , __a :							List[Any]=True       , __a :							Union[str, Any]=True       , __a :							Optional[Any]=True       , __a :							int=99       , __a :							Any=32       , __a :							Optional[Any]=5       , __a :							Tuple=4       , __a :							Union[str, Any]=37       , __a :							Any="gelu"       , __a :							Union[str, Any]=0.1       , __a :							Any=0.1       , __a :							List[str]=512       , __a :							Tuple=16       , __a :							int=2       , __a :							Optional[int]=0.02       , __a :							List[str]=4       , )		->					Any:
        _UpperCamelCase   :						Any    =							parent
        _UpperCamelCase   :						int    =							batch_size
        _UpperCamelCase   :						str    =							seq_length
        _UpperCamelCase   :						Any    =							is_training
        _UpperCamelCase   :						Tuple    =							use_attention_mask
        _UpperCamelCase   :						Optional[Any]    =							use_token_type_ids
        _UpperCamelCase   :						List[str]    =							use_labels
        _UpperCamelCase   :						List[str]    =							vocab_size
        _UpperCamelCase   :						Any    =							hidden_size
        _UpperCamelCase   :						str    =							num_hidden_layers
        _UpperCamelCase   :						Tuple    =							num_attention_heads
        _UpperCamelCase   :						Union[str, Any]    =							intermediate_size
        _UpperCamelCase   :						Dict    =							hidden_act
        _UpperCamelCase   :						str    =							hidden_dropout_prob
        _UpperCamelCase   :						Any    =							attention_probs_dropout_prob
        _UpperCamelCase   :						int    =							max_position_embeddings
        _UpperCamelCase   :						List[str]    =							type_vocab_size
        _UpperCamelCase   :						List[str]    =							type_sequence_label_size
        _UpperCamelCase   :						int    =							initializer_range
        _UpperCamelCase   :						Dict    =							num_choices

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					str:
        _UpperCamelCase   :						Tuple    =							ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size						)

        _UpperCamelCase   :						Optional[Any]    =							None
        if self.use_attention_mask:
              _UpperCamelCase   :						List[Any]    =							random_attention_mask([self.batch_size, self.seq_length]						)

        _UpperCamelCase   :						Optional[int]    =							None
        if self.use_token_type_ids:
              _UpperCamelCase   :						Optional[Any]    =							ids_tensor([self.batch_size, self.seq_length]       , self.type_vocab_size						)

        _UpperCamelCase   :						Optional[int]    =							RobertaConfig(
            vocab_size=self.vocab_size       , hidden_size=self.hidden_size       , num_hidden_layers=self.num_hidden_layers       , num_attention_heads=self.num_attention_heads       , intermediate_size=self.intermediate_size       , hidden_act=self.hidden_act       , hidden_dropout_prob=self.hidden_dropout_prob       , attention_probs_dropout_prob=self.attention_probs_dropout_prob       , max_position_embeddings=self.max_position_embeddings       , type_vocab_size=self.type_vocab_size       , is_decoder=__a       , initializer_range=self.initializer_range       , )

        return config, input_ids, token_type_ids, attention_mask

  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]						)		->					Optional[int]:
        _UpperCamelCase   :						str    =							self.prepare_config_and_inputs()
        _UpperCamelCase,	_UpperCamelCase,	_UpperCamelCase,	_UpperCamelCase   :						int    =							config_and_inputs
        _UpperCamelCase   :						List[Any]    =							{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
        return config, inputs_dict


  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					Any:
        _UpperCamelCase   :						List[str]    =							self.prepare_config_and_inputs()
        _UpperCamelCase,	_UpperCamelCase,	_UpperCamelCase,	_UpperCamelCase   :						Any    =							config_and_inputs

        _UpperCamelCase   :						Optional[int]    =							True
        _UpperCamelCase   :						str    =							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
        _UpperCamelCase   :						List[str]    =							ids_tensor([self.batch_size, self.seq_length]       , vocab_size=2						)

        return (
            config,
            input_ids,
            token_type_ids,
            encoder_hidden_states,
            encoder_attention_mask,
        )





@require_flax
class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase		,	unittest.TestCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Dict						  =       True

  SCREAMING_SNAKE_CASE__				:int						  =       (
      (
          FlaxRobertaModel,
          FlaxRobertaForCausalLM,
          FlaxRobertaForMaskedLM,
          FlaxRobertaForSequenceClassification,
          FlaxRobertaForTokenClassification,
          FlaxRobertaForMultipleChoice,
          FlaxRobertaForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					int:
        _UpperCamelCase   :						List[str]    =							FlaxRobertaModelTester(self						)


  @slow
  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					Tuple:
        for model_class_name in self.all_model_classes:
              _UpperCamelCase   :						int    =							model_class_name.from_pretrained("roberta-base"       , from_pt=__a						)
              _UpperCamelCase   :						Optional[Any]    =							model(np.ones((1, 1)						)						)
              self.assertIsNotNone(__a						)







"""simple docstring"""


import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
  import torch
  from torch import nn

  from transformers import (
      Adafactor,
      AdamW,
      get_constant_schedule,
      get_constant_schedule_with_warmup,
      get_cosine_schedule_with_warmup,
      get_cosine_with_hard_restarts_schedule_with_warmup,
      get_inverse_sqrt_schedule,
      get_linear_schedule_with_warmup,
      get_polynomial_decay_schedule_with_warmup,
  )




def        lowercase__     (       lowercase_		,lowercase_=10      )       ->		Tuple:




      """simple docstring"""

      _UpperCamelCase   :						str    =							[]
      for _ in range(lowercase_      ):
            lrs.append(scheduler.get_lr()[0]      )
            scheduler.step()
      return lrs




def        lowercase__     (       lowercase_		,lowercase_=10      )       ->		str:




      """simple docstring"""

      _UpperCamelCase   :						Optional[int]    =							[]
      for step in range(lowercase_      ):
            lrs.append(scheduler.get_lr()[0]      )
            scheduler.step()
            if step == num_steps // 2:
                  with tempfile.TemporaryDirectory() as tmpdirname:
                        _UpperCamelCase   :						List[str]    =							os.path.join(lowercase_		,"schedule.bin"      )
                        torch.save(scheduler.state_dict()		,lowercase_      )

                        _UpperCamelCase   :						Optional[int]    =							torch.load(lowercase_      )
                        scheduler.load_state_dict(lowercase_      )
      return lrs



@require_torch
class   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


  def  __SCREAMING_SNAKE_CASE					(		self :							Dict       , __a :							Dict       , __a :							Any       , __a :							List[str]						)		->					List[str]:
        self.assertEqual(len(__a						)       , len(__a						)						)
        for a, b in zip(__a       , __a						):
              self.assertAlmostEqual(__a       , __a       , delta=__a						)

  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					str:
        _UpperCamelCase   :						Optional[int]    =							torch.tensor([0.1, -0.2, -0.1]       , requires_grad=__a						)
        _UpperCamelCase   :						List[Any]    =							torch.tensor([0.4, 0.2, -0.5]						)
        _UpperCamelCase   :						Any    =							nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        _UpperCamelCase   :						Dict    =							AdamW(params=[w]       , lr=2e-1       , weight_decay=0.0						)
        for _ in range(100						):
              _UpperCamelCase   :						List[Any]    =							criterion(__a       , __a						)
              loss.backward()
              optimizer.step()
              w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
              w.grad.zero_()
        self.assertListAlmostEqual(w.tolist()       , [0.4, 0.2, -0.5]       , tol=1e-2						)


  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					Tuple:
        _UpperCamelCase   :						List[Any]    =							torch.tensor([0.1, -0.2, -0.1]       , requires_grad=__a						)
        _UpperCamelCase   :						Any    =							torch.tensor([0.4, 0.2, -0.5]						)
        _UpperCamelCase   :						List[str]    =							nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        _UpperCamelCase   :						Dict    =							Adafactor(
            params=[w]       , lr=1e-2       , eps=(1e-3_0, 1e-3)       , clip_threshold=1.0       , decay_rate=-0.8       , betaa=__a       , weight_decay=0.0       , relative_step=__a       , scale_parameter=__a       , warmup_init=__a       , )
        for _ in range(1000						):
              _UpperCamelCase   :						Optional[Any]    =							criterion(__a       , __a						)
              loss.backward()
              optimizer.step()
              w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
              w.grad.zero_()
        self.assertListAlmostEqual(w.tolist()       , [0.4, 0.2, -0.5]       , tol=1e-2						)





@require_torch
class   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Optional[Any]						  =       nn.Linear(50		,	50 ) if is_torch_available() else None
  SCREAMING_SNAKE_CASE__				:List[Any]						  =       AdamW(m.parameters()		,	lr=10.0 ) if is_torch_available() else None
  SCREAMING_SNAKE_CASE__				:List[str]						  =       10

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]       , __a :							Tuple       , __a :							List[Any]       , __a :							Dict       , __a :							str=None						)		->					List[Any]:
        self.assertEqual(len(__a						)       , len(__a						)						)
        for a, b in zip(__a       , __a						):
              self.assertAlmostEqual(__a       , __a       , delta=__a       , msg=__a						)


  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[Any]						)		->					List[Any]:
        _UpperCamelCase   :						Union[str, Any]    =							{"num_warmup_steps": 2, "num_training_steps": 10}
        # schedulers doct format
        # function: (sched_args_dict, expected_learning_rates)
        _UpperCamelCase   :						Optional[Any]    =							{
            get_constant_schedule: ({}, [10.0] * self.num_steps),
            get_constant_schedule_with_warmup: (
                {"num_warmup_steps": 4},
                [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
            ),
            get_linear_schedule_with_warmup: (
                {**common_kwargs},
                [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
            ),
            get_cosine_schedule_with_warmup: (
                {**common_kwargs},
                [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
            ),
            get_cosine_with_hard_restarts_schedule_with_warmup: (
                {**common_kwargs, "num_cycles": 2},
                [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
            ),
            get_polynomial_decay_schedule_with_warmup: (
                {**common_kwargs, "power": 2.0, "lr_end": 1e-7},
                [0.0, 5.0, 10.0, 7.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
            ),
            get_inverse_sqrt_schedule: (
                {"num_warmup_steps": 2},
                [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
            ),
        }

        for scheduler_func, data in scheds.items():
              _UpperCamelCase,	_UpperCamelCase   :						List[str]    =							data

              _UpperCamelCase   :						List[str]    =							scheduler_func(self.optimizer       , **__a						)
              self.assertEqual(len([scheduler.get_lr()[0]]						)       , 1						)
              _UpperCamelCase   :						Optional[int]    =							unwrap_schedule(__a       , self.num_steps						)
              self.assertListAlmostEqual(
                  __a       , __a       , tol=1e-2       , msg=F'''failed for {scheduler_func} in normal scheduler'''       , )

              _UpperCamelCase   :						List[str]    =							scheduler_func(self.optimizer       , **__a						)
              if scheduler_func.__name__ != "get_constant_schedule":
                    LambdaScheduleWrapper.wrap_scheduler(__a						)  # wrap to test picklability of the schedule
              _UpperCamelCase   :						int    =							unwrap_and_save_reload_schedule(__a       , self.num_steps						)
              self.assertListEqual(__a       , __a       , msg=F'''failed for {scheduler_func} in save and reload'''						)





class   __SCREAMING_SNAKE_CASE     :







  '''simple docstring'''


  def __init__(		self :							List[Any]       , __a :							Optional[int]						)		->					Tuple:
        _UpperCamelCase   :						Any    =							fn

  def __call__(		self :							Tuple       , *__a :							Optional[Any]       , **__a :							str						)		->					Any:
        return self.fn(*__a       , **__a						)


  @classmethod
  def  __SCREAMING_SNAKE_CASE					(		self :							Any       , __a :							Any						)		->					Tuple:
        _UpperCamelCase   :						Any    =							list(map(self       , scheduler.lr_lambdas						)						)
