






"""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 					_lowerCAmelCase(a   :    Union[str, Any]			,      a   :    str=10  )      ->       Optional[Any]:
	_SCREAMING_SNAKE_CASE						=[]
	for _ in range(_snake_case  ):
		lrs.append(scheduler.get_lr()[0]  )
		scheduler.step()
	return lrs






def 					_lowerCAmelCase(a   :    List[str]			,      a   :    Dict=10  )      ->       Optional[int]:
	_SCREAMING_SNAKE_CASE						=[]
	for step in range(_snake_case  ):
		lrs.append(scheduler.get_lr()[0]  )
		scheduler.step()
		if step == num_steps // 2:
			with tempfile.TemporaryDirectory() as tmpdirname:
				_SCREAMING_SNAKE_CASE						=os.path.join(_snake_case			,      '''schedule.bin'''  )
				torch.save(scheduler.state_dict()			,      _snake_case  )

				_SCREAMING_SNAKE_CASE						=torch.load(_snake_case  )
				scheduler.load_state_dict(_snake_case  )
	return lrs


@require_torch
class 	__UpperCAmelCase		(				unittest.TestCase	):


	'''simple docstring'''





	def 				UpperCamelCase_  (     self      ,      _A      ,      _A      ,      _A					):



		'''simple docstring'''

		self.assertEqual(len(__lowerCAmelCase					)      ,      len(__lowerCAmelCase					)					)
		for a, b in zip(__lowerCAmelCase      ,      __lowerCAmelCase					):
			self.assertAlmostEqual(__lowerCAmelCase      ,      __lowerCAmelCase      ,      delta=__lowerCAmelCase					)


	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=torch.tensor([0.1, -0.2, -0.1]      ,      requires_grad=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=torch.tensor([0.4, 0.2, -0.5]					)
		_SCREAMING_SNAKE_CASE						=nn.MSELoss()
		# No warmup, constant schedule, no gradient clipping
		_SCREAMING_SNAKE_CASE						=AdamW(params=[w]      ,      lr=2E-1      ,      weight_decay=0.0					)
		for _ in range(1_0_0					):
			_SCREAMING_SNAKE_CASE						=criterion(__lowerCAmelCase      ,      __lowerCAmelCase					)
			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					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=torch.tensor([0.1, -0.2, -0.1]      ,      requires_grad=__lowerCAmelCase					)
		_SCREAMING_SNAKE_CASE						=torch.tensor([0.4, 0.2, -0.5]					)
		_SCREAMING_SNAKE_CASE						=nn.MSELoss()
		# No warmup, constant schedule, no gradient clipping
		_SCREAMING_SNAKE_CASE						=Adafactor(
		    params=[w]      ,      lr=1E-2      ,      eps=(1E-30, 1E-3)      ,      clip_threshold=1.0      ,      decay_rate=-0.8      ,      betaa=__lowerCAmelCase      ,      weight_decay=0.0      ,      relative_step=__lowerCAmelCase      ,      scale_parameter=__lowerCAmelCase      ,      warmup_init=__lowerCAmelCase      ,      )
		for _ in range(1_0_0_0					):
			_SCREAMING_SNAKE_CASE						=criterion(__lowerCAmelCase      ,      __lowerCAmelCase					)
			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		:	Optional[Any]							     =				nn.Linear(50,  50	) if is_torch_available() else None
	lowercase		:	Dict							     =				AdamW(m.parameters(),  lr=1_0.0	) if is_torch_available() else None
	lowercase		:	List[Any]							     =				10


	def 				UpperCamelCase_  (     self      ,      _A      ,      _A      ,      _A      ,      _A=None					):



		'''simple docstring'''

		self.assertEqual(len(__lowerCAmelCase					)      ,      len(__lowerCAmelCase					)					)
		for a, b in zip(__lowerCAmelCase      ,      __lowerCAmelCase					):
			self.assertAlmostEqual(__lowerCAmelCase      ,      __lowerCAmelCase      ,      delta=__lowerCAmelCase      ,      msg=__lowerCAmelCase					)







	def 				UpperCamelCase_  (     self					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						={'''num_warmup_steps''': 2, '''num_training_steps''': 1_0}
		# schedulers doct format
		# function: (sched_args_dict, expected_learning_rates)
		_SCREAMING_SNAKE_CASE						={
		    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():
			_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE						=data

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

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




class 	__UpperCAmelCase		:


	'''simple docstring'''





	def __init__(     self      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=fn


	def __call__(     self      ,      *_A      ,      **_A					):



		'''simple docstring'''

		return self.fn(*__lowerCAmelCase      ,      **__lowerCAmelCase					)







	@classmethod
	def 				UpperCamelCase_  (     self      ,      _A					):



		'''simple docstring'''

		_SCREAMING_SNAKE_CASE						=list(map(self      ,      scheduler.lr_lambdas					)					)









import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_		 =					logging.get_logger(__name__)

UpperCAmelCase_		 =					{
    """xlnet-base-cased""": """https://huggingface.co/xlnet-base-cased/resolve/main/config.json""",
    """xlnet-large-cased""": """https://huggingface.co/xlnet-large-cased/resolve/main/config.json""",
}


class 	lowerCamelCase__	(       _A):



 """simple docstring"""

 a__					: Any =  "xlnet"
 a__					: Dict =  ["mems"]
 a__					: List[str] =  {
     "n_token": "vocab_size",  # Backward compatibility
     "hidden_size": "d_model",
     "num_attention_heads": "n_head",
     "num_hidden_layers": "n_layer",
 }

 def __init__(    self					:					int			,						__lowerCAmelCase					:					Dict=3_20_00			,						__lowerCAmelCase					:					List[str]=10_24			,						__lowerCAmelCase					:					Dict=24			,						__lowerCAmelCase					:					Optional[Any]=16			,						__lowerCAmelCase					:					Dict=40_96			,						__lowerCAmelCase					:					Any="gelu"			,						__lowerCAmelCase					:					int=True			,						__lowerCAmelCase					:					List[str]="bi"			,						__lowerCAmelCase					:					Dict=0.02			,						__lowerCAmelCase					:					Union[str, Any]=1E-12			,						__lowerCAmelCase					:					Optional[Any]=0.1			,						__lowerCAmelCase					:					Optional[Any]=5_12			,						__lowerCAmelCase					:					List[str]=None			,						__lowerCAmelCase					:					Tuple=True			,						__lowerCAmelCase					:					Tuple=False			,						__lowerCAmelCase					:					Optional[Any]=False			,						__lowerCAmelCase					:					Union[str, Any]=-1			,						__lowerCAmelCase					:					Optional[Any]=False			,						__lowerCAmelCase					:					Any="last"			,						__lowerCAmelCase					:					List[Any]=True			,						__lowerCAmelCase					:					Tuple="tanh"			,						__lowerCAmelCase					:					int=0.1			,						__lowerCAmelCase					:					str=5			,						__lowerCAmelCase					:					str=5			,						__lowerCAmelCase					:					List[str]=5			,						__lowerCAmelCase					:					List[str]=1			,						__lowerCAmelCase					:					Optional[int]=2			,						**__lowerCAmelCase					:					List[str]			,						)   ->	Tuple:
    _A      =       vocab_size
    _A      =       d_model
    _A      =       n_layer
    _A      =       n_head
    if d_model % n_head != 0:
       raise ValueError(f'''\'d_model % n_head\' ({d_model % n_head}) should be equal to 0'''      )
    if "d_head" in kwargs:
       if kwargs["d_head"] != d_model // n_head:
          raise ValueError(
              f'''`d_head` ({kwargs['d_head']}) should be equal to `d_model // n_head` ({d_model // n_head})'''      )
    _A      =       d_model // n_head
    _A      =       ff_activation
    _A      =       d_inner
    _A      =       untie_r
    _A      =       attn_type

    _A      =       initializer_range
    _A      =       layer_norm_eps

    _A      =       dropout
    _A      =       mem_len
    _A      =       reuse_len
    _A      =       bi_data
    _A      =       clamp_len
    _A      =       same_length

    _A      =       summary_type
    _A      =       summary_use_proj
    _A      =       summary_activation
    _A      =       summary_last_dropout
    _A      =       start_n_top
    _A      =       end_n_top

    _A      =       bos_token_id
    _A      =       pad_token_id
    _A      =       eos_token_id

    if "use_cache" in kwargs:
       warnings.warn(
           '''The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`'''
           ''' instead.'''			,						__lowerCAmelCase			,						)
       _A      =       kwargs['''use_cache''']

    _A      =       use_mems_eval
    _A      =       use_mems_train
    super().__init__(pad_token_id=__lowerCAmelCase			,						bos_token_id=__lowerCAmelCase			,						eos_token_id=__lowerCAmelCase			,						**__lowerCAmelCase      )

 @property
 def     snake_case_      (    self					:					Optional[Any]      )   ->	Union[str, Any]:
    logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''      )
    return -1

 @max_position_embeddings.setter
 def     snake_case_      (    self					:					Tuple			,						__lowerCAmelCase					:					Optional[Any]      )   ->	Dict:
    # Message copied from Transformer-XL documentation
    raise NotImplementedError(
        f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''      )



