





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


   A_     :				Union[str, Any]			      =	[]
   for _ in range(_lowerCAmelCase     ):
      lrs.append(scheduler.get_lr()[0]     )
      scheduler.step()
   return lrs



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase=1_0     ):
   '''simple docstring'''


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

            A_     :				Any			      =	torch.load(_lowerCAmelCase     )
            scheduler.load_state_dict(_lowerCAmelCase     )
   return lrs

@require_torch
class    _UpperCAmelCase							(		unittest.TestCase     ):

      '''simple docstring'''





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





      a						       =							nn.Linear(50			,       50     ) if is_torch_available() else None
      a						       =							AdamW(m.parameters()			,       lr=10.0     ) if is_torch_available() else None
      a						       =							10
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__      ,       a__=None       ):
         self.assertEqual(len(UpperCamelCase__       )      ,       len(UpperCamelCase__       )       )
         for a, b in zip(UpperCamelCase__      ,       UpperCamelCase__       ):
            self.assertAlmostEqual(UpperCamelCase__      ,       UpperCamelCase__      ,       delta=UpperCamelCase__      ,       msg=UpperCamelCase__       )






      def 		_lowerCamelCase      (		self       ):
         A_     :				Optional[Any]			      =	{"""num_warmup_steps""": 2, """num_training_steps""": 10}
         # schedulers doct format
         # function: (sched_args_dict, expected_learning_rates)
         A_     :				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():
            A_      ,				A_     :				Optional[Any]			      =	data

            A_     :				List[str]			      =	scheduler_func(self.optimizer      ,       **UpperCamelCase__       )
            self.assertEqual(len([scheduler.get_lr()[0]]       )      ,       1       )
            A_     :				Tuple			      =	unwrap_schedule(UpperCamelCase__      ,       self.num_steps       )
            self.assertListAlmostEqual(
                UpperCamelCase__      ,       UpperCamelCase__      ,       tol=1E-2      ,       msg=F"""failed for {scheduler_func} in normal scheduler"""      ,       )

            A_     :				str			      =	scheduler_func(self.optimizer      ,       **UpperCamelCase__       )
            if scheduler_func.__name__ != "get_constant_schedule":
               LambdaScheduleWrapper.wrap_scheduler(UpperCamelCase__       )  # wrap to test picklability of the schedule
            A_     :				Union[str, Any]			      =	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      ,       a__       ):
         A_     :				List[Any]			      =	fn
      def __call__(		self      ,       *a__      ,       **a__       ):
         return self.fn(*UpperCamelCase__      ,       **UpperCamelCase__       )






      @classmethod
      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Tuple			      =	list(map(self      ,       scheduler.lr_lambdas       )       )













from typing import List

from .keymap import KEYMAP, get_character



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
   '''simple docstring'''


   def decorator(_lowerCAmelCase     ):
      A_     :				List[Any]			      =	getattr(_lowerCAmelCase  ,"""handle_key"""  ,[]     )
      handle += [key]
      setattr(_lowerCAmelCase  ,"""handle_key"""  ,_lowerCAmelCase     )
      return func

   return decorator



def     _lowerCAmelCase		(  *_lowerCAmelCase     ):
   '''simple docstring'''


   def decorator(_lowerCAmelCase     ):
      A_     :				Tuple			      =	getattr(_lowerCAmelCase  ,"""handle_key"""  ,[]     )
      handle += keys
      setattr(_lowerCAmelCase  ,"""handle_key"""  ,_lowerCAmelCase     )
      return func

   return decorator

class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __new__(		cls      ,       a__      ,       a__      ,       a__       ):
         A_     :				Any			      =	super().__new__(cls      ,       a__      ,       a__      ,       a__       )
         if not hasattr(a__      ,       """key_handler"""       ):
            setattr(a__      ,       """key_handler"""      ,       {}       )
         setattr(a__      ,       """handle_input"""      ,       KeyHandler.handle_input       )

         for value in attrs.values():
            A_     :				Tuple			      =	getattr(a__      ,       """handle_key"""      ,       []       )
            for key in handled_keys:
               A_     :				Optional[Any]			      =	value
         return new_cls






      @staticmethod
      def 		_lowerCamelCase      (		cls       ):
         A_     :				List[str]			      =	get_character()
         if char != KEYMAP["undefined"]:
            A_     :				str			      =	ord(a__       )
         A_     :				List[str]			      =	cls.key_handler.get(a__       )
         if handler:
            A_     :				Optional[int]			      =	char
            return handler(cls       )
         else:
            return None






def     _lowerCAmelCase		(  cls     ):
   '''simple docstring'''


   return KeyHandler(cls.__name__  ,cls.__bases__  ,cls.__dict__.copy()     )






