


"""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       A_   (   _lowerCAmelCase     :					Dict,     _lowerCAmelCase     :					Optional[Any]=10       ):






     """simple docstring"""

     _a    	=				[]
     for _ in range(_lowerCAmelCase       ):
          lrs.append(scheduler.get_lr()[0]       )
          scheduler.step()
     return lrs
def       A_   (   _lowerCAmelCase     :					Optional[int],     _lowerCAmelCase     :					List[str]=10       ):






     """simple docstring"""

     _a    	=				[]
     for step in range(_lowerCAmelCase       ):
          lrs.append(scheduler.get_lr()[0]       )
          scheduler.step()
          if step == num_steps // 2:
               with tempfile.TemporaryDirectory() as tmpdirname:
                    _a    	=				os.path.join(_lowerCAmelCase,     '''schedule.bin'''       )
                    torch.save(scheduler.state_dict(),     _lowerCAmelCase       )

                    _a    	=				torch.load(_lowerCAmelCase       )
                    scheduler.load_state_dict(_lowerCAmelCase       )
     return lrs


@require_torch
class     __lowerCamelCase  (					unittest.TestCase				):

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							List[str]:
          self.assertEqual(len(__UpperCAmelCase							)    ,				len(__UpperCAmelCase							)							)
          for a, b in zip(__UpperCAmelCase    ,				__UpperCAmelCase							):
               self.assertAlmostEqual(__UpperCAmelCase    ,				__UpperCAmelCase    ,				delta=__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self							)       ->							Optional[Any]:
          _a    	=				torch.tensor([0.1, -0.2, -0.1]    ,				requires_grad=__UpperCAmelCase							)
          _a    	=				torch.tensor([0.4, 0.2, -0.5]							)
          _a    	=				nn.MSELoss()
          # No warmup, constant schedule, no gradient clipping
          _a    	=				AdamW(params=[w]    ,				lr=2e-1    ,				weight_decay=0.0							)
          for _ in range(100							):
               _a    	=				criterion(__UpperCAmelCase    ,				__UpperCAmelCase							)
               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 							_UpperCAmelCase      (      self							)       ->							str:
          _a    	=				torch.tensor([0.1, -0.2, -0.1]    ,				requires_grad=__UpperCAmelCase							)
          _a    	=				torch.tensor([0.4, 0.2, -0.5]							)
          _a    	=				nn.MSELoss()
          # No warmup, constant schedule, no gradient clipping
          _a    	=				Adafactor(
              params=[w]    ,				lr=1e-2    ,				eps=(1e-3_0, 1e-3)    ,				clip_threshold=1.0    ,				decay_rate=-0.8    ,				betaa=__UpperCAmelCase    ,				weight_decay=0.0    ,				relative_step=__UpperCAmelCase    ,				scale_parameter=__UpperCAmelCase    ,				warmup_init=__UpperCAmelCase    ,				)
          for _ in range(1000							):
               _a    	=				criterion(__UpperCAmelCase    ,				__UpperCAmelCase							)
               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     __lowerCamelCase  (					unittest.TestCase				):

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


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None							)       ->							int:
          self.assertEqual(len(__UpperCAmelCase							)    ,				len(__UpperCAmelCase							)							)
          for a, b in zip(__UpperCAmelCase    ,				__UpperCAmelCase							):
               self.assertAlmostEqual(__UpperCAmelCase    ,				__UpperCAmelCase    ,				delta=__UpperCAmelCase    ,				msg=__UpperCAmelCase							)





     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          _a    	=				{'''num_warmup_steps''': 2, '''num_training_steps''': 10}
          # schedulers doct format
          # function: (sched_args_dict, expected_learning_rates)
          _a    	=				{
              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.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
              ),
              get_inverse_sqrt_schedule: (
                  {'''num_warmup_steps''': 2},
                  [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
              ),
          }

          for scheduler_func, data in scheds.items():
               _a	,      _a    	=				data

               _a    	=				scheduler_func(self.optimizer    ,				**__UpperCAmelCase							)
               self.assertEqual(len([scheduler.get_lr()[0]]							)    ,				1							)
               _a    	=				unwrap_schedule(__UpperCAmelCase    ,				self.num_steps							)
               self.assertListAlmostEqual(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				tol=1e-2    ,				msg=F'failed for {scheduler_func} in normal scheduler'    ,				)

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







class     __lowerCamelCase  :

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase							)       ->							str:
          _a    	=				fn


     def __call__(      self    ,				*__UpperCAmelCase    ,				**__UpperCAmelCase							)       ->							Any:
          return self.fn(*__UpperCAmelCase    ,				**__UpperCAmelCase							)





     @classmethod
     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          _a    	=				list(map(self    ,				scheduler.lr_lambdas							)							)



"""simple docstring"""






import argparse

import torch

from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
from transformers.utils import logging


logging.set_verbosity_info()
def       A_   (   _lowerCAmelCase     :					Dict,     _lowerCAmelCase     :					List[Any],     _lowerCAmelCase     :					Union[str, Any]       ):






     """simple docstring"""

     _a    	=				BertConfig.from_json_file(_lowerCAmelCase       )
     print(f'Building PyTorch model from configuration: {config}'       )
     _a    	=				BertForPreTraining(_lowerCAmelCase       )

     # Load weights from tf checkpoint
     load_tf_weights_in_bert(_lowerCAmelCase,     _lowerCAmelCase,     _lowerCAmelCase       )

     # Save pytorch-model
     print(f'Save PyTorch model to {pytorch_dump_path}'       )
     torch.save(model.state_dict(),     _lowerCAmelCase       )


if __name__ == "__main__":
   __snake_case		=       argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
   )
   parser.add_argument(
       '''--bert_config_file''',
       default=None,
       type=str,
       required=True,
       help=(
           '''The config json file corresponding to the pre-trained BERT model. \n'''
           '''This specifies the model architecture.'''
       ),
   )
   parser.add_argument(
       '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
   )
   __snake_case		=       parser.parse_args()
   convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)