


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     UpperCamelCase    (      _a       ,		_a=1_0		)					->							int:
     '''simple docstring'''





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




def     UpperCamelCase    (      _a       ,		_a=1_0		)					->							List[Any]:
     '''simple docstring'''





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

                    lowercase_       :Optional[Any]     =       torch.load(_a		)
                    scheduler.load_state_dict(_a		)
     return lrs






@require_torch
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        self.assertEqual(len(UpperCamelCase_		)   ,							len(UpperCamelCase_		)		)
        for a, b in zip(UpperCamelCase_   ,							UpperCamelCase_		):
             self.assertAlmostEqual(UpperCamelCase_   ,							UpperCamelCase_   ,							delta=UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       torch.tensor([0.1, -0.2, -0.1]   ,							requires_grad=UpperCamelCase_		)
        lowercase_       :Any     =       torch.tensor([0.4, 0.2, -0.5]		)
        lowercase_       :Tuple     =       nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        lowercase_       :str     =       AdamW(params=[w]   ,							lr=2E-1   ,							weight_decay=0.0		)
        for _ in range(100		):
             lowercase_       :List[str]     =       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		):
        lowercase_       :Optional[int]     =       torch.tensor([0.1, -0.2, -0.1]   ,							requires_grad=UpperCamelCase_		)
        lowercase_       :str     =       torch.tensor([0.4, 0.2, -0.5]		)
        lowercase_       :List[Any]     =       nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        lowercase_       :Optional[Any]     =       Adafactor(
            params=[w]   ,							lr=1E-2   ,							eps=(1E-30, 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		):
             lowercase_       :List[str]     =       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      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   lowercase					:	List[Any]									=nn.Linear(50						,      50     ) if is_torch_available() else None
   lowercase					:	List[Any]									=AdamW(m.parameters()						,      lr=10.0     ) if is_torch_available() else None
   lowercase					:	List[str]									=10
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=None		):
        self.assertEqual(len(UpperCamelCase_		)   ,							len(UpperCamelCase_		)		)
        for a, b in zip(UpperCamelCase_   ,							UpperCamelCase_		):
             self.assertAlmostEqual(UpperCamelCase_   ,							UpperCamelCase_   ,							delta=UpperCamelCase_   ,							msg=UpperCamelCase_		)




   def 		UpperCamelCase  (     self		):
        lowercase_       :Tuple     =       {'''num_warmup_steps''': 2, '''num_training_steps''': 10}
        # schedulers doct format
        # function: (sched_args_dict, expected_learning_rates)
        lowercase_       :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.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():
             lowercase_    ,      lowercase_       :int     =       data

             lowercase_       :Any     =       scheduler_func(self.optimizer   ,							**UpperCamelCase_		)
             self.assertEqual(len([scheduler.get_lr()[0]]		)   ,							1		)
             lowercase_       :Optional[int]     =       unwrap_schedule(UpperCamelCase_   ,							self.num_steps		)
             self.assertListAlmostEqual(
                 UpperCamelCase_   ,							UpperCamelCase_   ,							tol=1E-2   ,							msg=f"failed for {scheduler_func} in normal scheduler"   ,							)

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






class      UpperCamelCase					:


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_		):
        lowercase_       :Dict     =       fn
   def __call__(     self   ,							*UpperCamelCase_   ,							**UpperCamelCase_		):
        return self.fn(*UpperCamelCase_   ,							**UpperCamelCase_		)




   @classmethod
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        lowercase_       :Union[str, Any]     =       list(map(self   ,							scheduler.lr_lambdas		)		)





import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()




def     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							List[str]:
     '''simple docstring'''





     if openai_config_file == "":
          lowercase_       :str     =       OpenAIGPTConfig()
     else:
          lowercase_       :int     =       OpenAIGPTConfig.from_json_file(_a		)
     lowercase_       :int     =       OpenAIGPTModel(_a		)

     # Load weights from numpy
     load_tf_weights_in_openai_gpt(_a       ,		_a       ,		_a		)

     # Save pytorch-model
     lowercase_       :Optional[int]     =       pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
     lowercase_       :List[str]     =       pytorch_dump_folder_path + '''/''' + CONFIG_NAME
     print(f"Save PyTorch model to {pytorch_weights_dump_path}"		)
     torch.save(model.state_dict()       ,		_a		)
     print(f"Save configuration file to {pytorch_config_dump_path}"		)
     with open(_a       ,		'''w'''       ,		encoding='''utf-8'''		) as f:
          f.write(config.to_json_string()		)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE       :		List[str]         =				argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       "--openai_checkpoint_folder_path",
       default=None,
       type=str,
       required=True,
       help="Path to the TensorFlow checkpoint path.",
   )
   parser.add_argument(
       "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
   )
   parser.add_argument(
       "--openai_config_file",
       default="",
       type=str,
       help=(
           "An optional config json file corresponding to the pre-trained OpenAI model. \n"
           "This specifies the model architecture."
       ),
   )
   SCREAMING_SNAKE_CASE       :		Optional[Any]         =				parser.parse_args()
   convert_openai_checkpoint_to_pytorch(
       args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
   )

