




import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
  )


class 							UpperCamelCase_       :

    '''simple docstring'''






    def __init__(     self :      List[Any]   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      Tuple=3   ,					UpperCAmelCase__ :      List[Any]=7   ,					UpperCAmelCase__ :      List[str]=True   ,					UpperCAmelCase__ :      Any=True   ,					UpperCAmelCase__ :      Union[str, Any]=False   ,					UpperCAmelCase__ :      Dict=True   ,					UpperCAmelCase__ :      Any=99   ,					UpperCAmelCase__ :      Optional[int]=32   ,					UpperCAmelCase__ :      str=5   ,					UpperCAmelCase__ :      Dict=4   ,					UpperCAmelCase__ :      str=37   ,					UpperCAmelCase__ :      int="gelu"   ,					UpperCAmelCase__ :      Dict=0.1   ,					UpperCAmelCase__ :      List[str]=0.1   ,					UpperCAmelCase__ :      Any=512   ,					UpperCAmelCase__ :      Union[str, Any]=16   ,					UpperCAmelCase__ :      Tuple=2   ,					UpperCAmelCase__ :      Any=0.02   ,					UpperCAmelCase__ :      List[Any]=3   ,					UpperCAmelCase__ :      Tuple=4   ,					UpperCAmelCase__ :      Union[str, Any]=None   ,					)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       parent
           A__ 					=       batch_size
           A__ 					=       seq_length
           A__ 					=       is_training
           A__ 					=       use_input_mask
           A__ 					=       use_token_type_ids
           A__ 					=       use_labels
           A__ 					=       vocab_size
           A__ 					=       hidden_size
           A__ 					=       num_hidden_layers
           A__ 					=       num_attention_heads
           A__ 					=       intermediate_size
           A__ 					=       hidden_act
           A__ 					=       hidden_dropout_prob
           A__ 					=       attention_probs_dropout_prob
           A__ 					=       max_position_embeddings
           A__ 					=       type_vocab_size
           A__ 					=       type_sequence_label_size
           A__ 					=       initializer_range
           A__ 					=       num_labels
           A__ 					=       num_choices
           A__ 					=       scope


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Tuple:





           '''simple docstring'''

           A__ 					=       ids_tensor([self.batch_size, self.seq_length]   ,					self.vocab_size)

           A__ 					=       None
           if self.use_input_mask:
                  A__ 					=       random_attention_mask([self.batch_size, self.seq_length])

           A__ 					=       None

           A__ 					=       None
           A__ 					=       None
           A__ 					=       None
           if self.use_labels:
                  A__ 					=       ids_tensor([self.batch_size]   ,					self.type_sequence_label_size)
                  A__ 					=       ids_tensor([self.batch_size, self.seq_length]   ,					self.num_labels)
                  A__ 					=       ids_tensor([self.batch_size]   ,					self.num_choices)

           A__ 					=       self.get_config()

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


    def 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any])		->Optional[int]:





           '''simple docstring'''

           return FalconConfig(
               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=UpperCAmelCase__   ,					initializer_range=self.initializer_range   ,					pad_token_id=1   ,					new_decoder_architecture=UpperCAmelCase__   ,					)


    def 			SCREAMING_SNAKE_CASE				(     self :      Any   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      List[str])		->List[str]:





           '''simple docstring'''

           A__ 					=       FalconModel(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__)
           A__ 					=       model(UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape   ,					(self.batch_size, self.seq_length, self.hidden_size))


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Optional[int]   ,					)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       True
           A__ 					=       FalconModel(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(
               UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					encoder_hidden_states=UpperCAmelCase__   ,					encoder_attention_mask=UpperCAmelCase__   ,					)
           A__ 					=       model(
               UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					encoder_hidden_states=UpperCAmelCase__   ,					)
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape   ,					(self.batch_size, self.seq_length, self.hidden_size))


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					)		->int:





           '''simple docstring'''

           A__ 					=       FalconForCausalLM(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					labels=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.seq_length, self.vocab_size))


    def 			SCREAMING_SNAKE_CASE				(     self :      Any   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      int   ,					)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       True
           A__ 					=       True
           A__ 					=       FalconForCausalLM(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()

           # first forward pass
           A__ 					=       model(
               UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					encoder_hidden_states=UpperCAmelCase__   ,					encoder_attention_mask=UpperCAmelCase__   ,					use_cache=UpperCAmelCase__   ,					)
           A__ 					=       outputs.past_key_values

           # create hypothetical multiple next token and extent to next_input_ids
           A__ 					=       ids_tensor((self.batch_size, 3)   ,					config.vocab_size)
           A__ 					=       ids_tensor((self.batch_size, 3)   ,					vocab_size=2)

           # append to next input_ids and
           A__ 					=       torch.cat([input_ids, next_tokens]   ,					dim=-1)
           A__ 					=       torch.cat([input_mask, next_mask]   ,					dim=-1)

           A__ 					=       model(
               UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					encoder_hidden_states=UpperCAmelCase__   ,					encoder_attention_mask=UpperCAmelCase__   ,					output_hidden_states=UpperCAmelCase__   ,					)['''hidden_states'''][0]
           A__ 					=       model(
               UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					encoder_hidden_states=UpperCAmelCase__   ,					encoder_attention_mask=UpperCAmelCase__   ,					past_key_values=UpperCAmelCase__   ,					output_hidden_states=UpperCAmelCase__   ,					)['''hidden_states'''][0]

           # select random slice
           A__ 					=       ids_tensor((1,)   ,					output_from_past.shape[-1]).item()
           A__ 					=       output_from_no_past[:, -3:, random_slice_idx].detach()
           A__ 					=       output_from_past[:, :, random_slice_idx].detach()

           self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

           # test that outputs are equal for slice
           self.parent.assertTrue(torch.allclose(UpperCAmelCase__   ,					UpperCAmelCase__   ,					atol=1e-3))





    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->List[str]:





           '''simple docstring'''

           A__ 					=       self.prepare_config_and_inputs()
           (
               (
               A__
           )					,				(
               A__
           )					,				(
               A__
           )					,				(
               A__
           )					,				(
               A__
           )					,				(
               A__
           )					,				(
               A__
           )					,				
           ) 					=       config_and_inputs
           A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict


@require_torch
class 							UpperCamelCase_       (					UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = (
        (
            FalconModel,
            FalconForCausalLM,
            FalconForSequenceClassification,
            FalconForTokenClassification,
            FalconForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    UpperCAmelCase__				    = (FalconForCausalLM,) if is_torch_available() else ()
    UpperCAmelCase__				    = (
        {
            '''feature-extraction''': FalconModel,
            '''text-classification''': FalconForSequenceClassification,
            '''text-generation''': FalconForCausalLM,
            '''question-answering''': FalconForQuestionAnswering,
            '''token-classification''': FalconForTokenClassification,
            '''zero-shot''': FalconForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->Tuple:





           '''simple docstring'''

           A__ 					=       FalconModelTester(self)
           A__ 					=       ConfigTester(self   ,					config_class=UpperCAmelCase__   ,					hidden_size=37)


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Union[str, Any]:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->Union[str, Any]:





           '''simple docstring'''

           A__					,				*A__ 					=       self.model_tester.prepare_config_and_inputs()
           for alibi in [True, False]:
                  A__ 					=       alibi
                  self.model_tester.create_and_check_model(UpperCAmelCase__   ,					*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->Optional[Any]:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()
           A__ 					=       3
           A__ 					=       input_dict['''input_ids''']
           A__ 					=       input_ids.ne(1).to(UpperCAmelCase__)
           A__ 					=       ids_tensor([self.model_tester.batch_size]   ,					self.model_tester.type_sequence_label_size)
           A__ 					=       FalconForSequenceClassification(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					labels=UpperCAmelCase__)
           self.assertEqual(result.logits.shape   ,					(self.model_tester.batch_size, self.model_tester.num_labels))


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->List[Any]:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()
           A__ 					=       3
           A__ 					=       '''single_label_classification'''
           A__ 					=       input_dict['''input_ids''']
           A__ 					=       input_ids.ne(1).to(UpperCAmelCase__)
           A__ 					=       ids_tensor([self.model_tester.batch_size]   ,					self.model_tester.type_sequence_label_size)
           A__ 					=       FalconForSequenceClassification(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					labels=UpperCAmelCase__)
           self.assertEqual(result.logits.shape   ,					(self.model_tester.batch_size, self.model_tester.num_labels))


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Dict:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()
           A__ 					=       input_dict['''input_ids''']
           A__ 					=       FalconForCausalLM(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					use_cache=UpperCAmelCase__)
           A__ 					=       input_ids.shape[0]
           A__ 					=       model._convert_to_rw_cache(result.past_key_values)
           A__ 					=       model._convert_cache_to_standard_format(UpperCAmelCase__   ,					UpperCAmelCase__)
           for layer in range(len(UpperCAmelCase__)):
                  for tensor_idx in range(2):
                         self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3)
                         self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4)
                         self.assertTrue(
                             torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]))


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->List[str]:





           '''simple docstring'''

           A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()
           A__ 					=       3
           A__ 					=       '''multi_label_classification'''
           A__ 					=       input_dict['''input_ids''']
           A__ 					=       input_ids.ne(1).to(UpperCAmelCase__)
           A__ 					=       ids_tensor(
               [self.model_tester.batch_size, config.num_labels]   ,					self.model_tester.type_sequence_label_size).to(torch.float)
           A__ 					=       FalconForSequenceClassification(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__ 					=       model(UpperCAmelCase__   ,					attention_mask=UpperCAmelCase__   ,					labels=UpperCAmelCase__)
           self.assertEqual(result.logits.shape   ,					(self.model_tester.batch_size, self.model_tester.num_labels))





    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->List[str]:





           '''simple docstring'''

           for model_class in self.all_generative_model_classes:
                  A__					,				A__ 					=       self.model_tester.prepare_config_and_inputs_for_common()

                  # If it doesn't support cache, pass the test
                  if not hasattr(UpperCAmelCase__   ,					'''use_cache'''):
                         return

                  A__ 					=       model_class(UpperCAmelCase__).to(UpperCAmelCase__)
                  if "use_cache" not in inputs:
                         A__ 					=       True
                  A__ 					=       model(**UpperCAmelCase__)

                  # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
                  if "past_key_values" not in outputs:
                         return

                  A__ 					=       (
                      getattr(UpperCAmelCase__   ,					'''decoder_layers'''   ,					UpperCAmelCase__)
                      or getattr(UpperCAmelCase__   ,					'''num_decoder_layers'''   ,					UpperCAmelCase__)
                      or config.num_hidden_layers
                  )
                  A__ 					=       getattr(UpperCAmelCase__   ,					'''num_kv_heads'''   ,					config.num_attention_heads)
                  A__ 					=       getattr(UpperCAmelCase__   ,					'''d_model'''   ,					config.hidden_size)
                  A__ 					=       embed_dim // num_attention_heads

                  A__ 					=       outputs['''past_key_values''']
                  self.assertEqual(len(UpperCAmelCase__)   ,					UpperCAmelCase__)

                  A__					,				A__ 					=       inputs['''input_ids'''].shape
                  for i in range(UpperCAmelCase__):
                         if config.new_decoder_architecture:
                                A__ 					=       config.num_attention_heads
                         elif config.multi_query:
                                A__ 					=       1
                         self.assertEqual(len(past_kv[0])   ,					2)  # K V for the decoder = 2
                         self.assertEqual(
                             past_kv[i][0].shape   ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim))
                         self.assertEqual(
                             past_kv[i][1].shape   ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim))


@require_torch
class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b''')
           A__ 					=       FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b''')
           model.eval()
           model.to(UpperCAmelCase__)
           A__ 					=       tokenizer('''My favorite food is'''   ,					return_tensors='''pt''').to(UpperCAmelCase__)

           A__ 					=       (
               '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
           )

           A__ 					=       model.generate(**UpperCAmelCase__   ,					do_sample=UpperCAmelCase__   ,					max_new_tokens=19)
           A__ 					=       tokenizer.batch_decode(UpperCAmelCase__)[0]

           self.assertEqual(UpperCAmelCase__   ,					UpperCAmelCase__)


    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->List[Any]:





           '''simple docstring'''

           for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
                  A__ 					=       AutoTokenizer.from_pretrained(UpperCAmelCase__)
                  A__ 					=       FalconForCausalLM.from_pretrained(UpperCAmelCase__)
                  model.eval()
                  model.to(UpperCAmelCase__)
                  A__ 					=       tokenizer('''My favorite food is'''   ,					return_tensors='''pt''').to(UpperCAmelCase__)

                  # We just test that these run without errors - the models are randomly initialized
                  # and so the actual text outputs will be garbage
                  model.generate(**UpperCAmelCase__   ,					do_sample=UpperCAmelCase__   ,					max_new_tokens=4)
                  model.generate(**UpperCAmelCase__   ,					do_sample=UpperCAmelCase__   ,					max_new_tokens=4)
                  model.generate(**UpperCAmelCase__   ,					num_beams=2   ,					max_new_tokens=4)





    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->int:





           '''simple docstring'''

           with torch.no_grad():
                  for repo in [
                      "Rocketknight1/falcon-rw-1b",
                      "Rocketknight1/tiny-random-falcon-7b",
                      "Rocketknight1/tiny-random-falcon-40b",
                  ]:
                         A__ 					=       AutoTokenizer.from_pretrained(UpperCAmelCase__)
                         A__ 					=       FalconForCausalLM.from_pretrained(UpperCAmelCase__)
                         model.eval()
                         model.to(device=UpperCAmelCase__)
                         A__ 					=       tokenizer('''My favorite food is'''   ,					return_tensors='''pt''').to(UpperCAmelCase__)

                         # Test results are the same with and without cache
                         A__ 					=       model.generate(**UpperCAmelCase__   ,					do_sample=UpperCAmelCase__   ,					max_new_tokens=20   ,					use_cache=UpperCAmelCase__)
                         A__ 					=       model.generate(**UpperCAmelCase__   ,					do_sample=UpperCAmelCase__   ,					max_new_tokens=20   ,					use_cache=UpperCAmelCase__)
                         self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0)








import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


a__		     =     logging.get_logger(__name__)


class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



	"""simple docstring"""


	__magic_name__   :       Dict      	=		'linear'
	__magic_name__   :       Dict      	=		'cosine'
	__magic_name__   :       Optional[int]      	=		'cosine_with_restarts'
	__magic_name__   :       List[str]      	=		'polynomial'
	__magic_name__   :       Any      	=		'constant'
	__magic_name__   :       Union[str, Any]      	=		'constant_with_warmup'
	__magic_name__   :       str      	=		'piecewise_constant'






def 				A__ (snake_case   :		Optimizer		,	snake_case   :		int = -1				)     ->						Optional[Any]:
					return LambdaLR(snake_case		,	lambda snake_case				: 1		,	last_epoch=snake_case				)





def 				A__ (snake_case   :		Optimizer		,	snake_case   :		int		,	snake_case   :		int = -1				)     ->						List[Any]:

					def lr_lambda(snake_case   :		int				):
										if current_step < num_warmup_steps:
															return float(snake_case				) / float(max(1.0		,	snake_case				)				)
										return 1.0

					return LambdaLR(snake_case		,	snake_case		,	last_epoch=snake_case				)





def 				A__ (snake_case   :		Optimizer		,	snake_case   :		str		,	snake_case   :		int = -1				)     ->						Union[str, Any]:
					__UpperCamelCase					:	Optional[Any]             =      {}
					__UpperCamelCase					:	int             =      step_rules.split(""","""				)
					for rule_str in rule_list[:-1]:
										__UpperCamelCase      ,			__UpperCamelCase					:	Tuple             =      rule_str.split(""":"""				)
										__UpperCamelCase					:	int             =      int(snake_case				)
										__UpperCamelCase					:	Union[str, Any]             =      float(snake_case				)
										__UpperCamelCase					:	Optional[int]             =      value
					__UpperCamelCase					:	Dict             =      float(rule_list[-1]				)

					def create_rules_function(snake_case   :		List[str]		,	snake_case   :		Any				):
										def rule_func(snake_case   :		int				) -> float:
															__UpperCamelCase					:	Any             =      sorted(rules_dict.keys()				)
															for i, sorted_step in enumerate(snake_case				):
																				if steps < sorted_step:
																									return rules_dict[sorted_steps[i]]
															return last_lr_multiple

										return rule_func

					__UpperCamelCase					:	Tuple             =      create_rules_function(snake_case		,	snake_case				)

					return LambdaLR(snake_case		,	snake_case		,	last_epoch=snake_case				)





def 				A__ (snake_case   :		int		,	snake_case   :		Optional[Any]		,	snake_case   :		Union[str, Any]		,	snake_case   :		str=-1				)     ->						str:

					def lr_lambda(snake_case   :		int				):
										if current_step < num_warmup_steps:
															return float(snake_case				) / float(max(1		,	snake_case				)				)
										return max(
										    0.0		,	float(num_training_steps - current_step				) / float(max(1		,	num_training_steps - num_warmup_steps				)				)				)

					return LambdaLR(snake_case		,	snake_case		,	snake_case				)





def 				A__ (snake_case   :		Optimizer		,	snake_case   :		int		,	snake_case   :		int		,	snake_case   :		float = 0.5		,	snake_case   :		int = -1				)     ->						List[str]:

					def lr_lambda(snake_case   :		Dict				):
										if current_step < num_warmup_steps:
															return float(snake_case				) / float(max(1		,	snake_case				)				)
										__UpperCamelCase					:	Dict             =      float(current_step - num_warmup_steps				) / float(max(1		,	num_training_steps - num_warmup_steps				)				)
										return max(0.0		,	0.5 * (1.0 + math.cos(math.pi * float(snake_case				) * 2.0 * progress				))				)

					return LambdaLR(snake_case		,	snake_case		,	snake_case				)





def 				A__ (snake_case   :		Optimizer		,	snake_case   :		int		,	snake_case   :		int		,	snake_case   :		int = 1		,	snake_case   :		int = -1				)     ->						Tuple:

					def lr_lambda(snake_case   :		Optional[int]				):
										if current_step < num_warmup_steps:
															return float(snake_case				) / float(max(1		,	snake_case				)				)
										__UpperCamelCase					:	List[str]             =      float(current_step - num_warmup_steps				) / float(max(1		,	num_training_steps - num_warmup_steps				)				)
										if progress >= 1.0:
															return 0.0
										return max(0.0		,	0.5 * (1.0 + math.cos(math.pi * ((float(snake_case				) * progress) % 1.0)				))				)

					return LambdaLR(snake_case		,	snake_case		,	snake_case				)





def 				A__ (snake_case   :		List[str]		,	snake_case   :		Optional[Any]		,	snake_case   :		Optional[int]		,	snake_case   :		str=1e-7		,	snake_case   :		List[str]=1.0		,	snake_case   :		Dict=-1				)     ->						Tuple:
					__UpperCamelCase					:	Tuple             =      optimizer.defaults["""lr"""]
					if not (lr_init > lr_end):
										raise ValueError(F'''lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})'''				)

					def lr_lambda(snake_case   :		int				):
										if current_step < num_warmup_steps:
															return float(snake_case				) / float(max(1		,	snake_case				)				)
										elif current_step > num_training_steps:
															return lr_end / lr_init  # as LambdaLR multiplies by lr_init
										else:
															__UpperCamelCase					:	List[str]             =      lr_init - lr_end
															__UpperCamelCase					:	Any             =      num_training_steps - num_warmup_steps
															__UpperCamelCase					:	List[str]             =      1 - (current_step - num_warmup_steps) / decay_steps
															__UpperCamelCase					:	List[Any]             =      lr_range * pct_remaining**power + lr_end
															return decay / lr_init  # as LambdaLR multiplies by lr_init

					return LambdaLR(snake_case		,	snake_case		,	snake_case				)


a__		     =     {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}







def 				A__ (snake_case   :		Union[str, SchedulerType]		,	snake_case   :		Optimizer		,	snake_case   :		Optional[str] = None		,	snake_case   :		Optional[int] = None		,	snake_case   :		Optional[int] = None		,	snake_case   :		int = 1		,	snake_case   :		float = 1.0		,	snake_case   :		int = -1		,	)     ->						Dict:
					__UpperCamelCase					:	List[str]             =      SchedulerType(snake_case				)
					__UpperCamelCase					:	int             =      TYPE_TO_SCHEDULER_FUNCTION[name]
					if name == SchedulerType.CONSTANT:
										return schedule_func(snake_case		,	last_epoch=snake_case				)

					if name == SchedulerType.PIECEWISE_CONSTANT:
										return schedule_func(snake_case		,	step_rules=snake_case		,	last_epoch=snake_case				)

					# All other schedulers require `num_warmup_steps`
					if num_warmup_steps is None:
										raise ValueError(F'''{name} requires `num_warmup_steps`, please provide that argument.'''				)

					if name == SchedulerType.CONSTANT_WITH_WARMUP:
										return schedule_func(snake_case		,	num_warmup_steps=snake_case		,	last_epoch=snake_case				)

					# All other schedulers require `num_training_steps`
					if num_training_steps is None:
										raise ValueError(F'''{name} requires `num_training_steps`, please provide that argument.'''				)

					if name == SchedulerType.COSINE_WITH_RESTARTS:
										return schedule_func(
										    snake_case		,	num_warmup_steps=snake_case		,	num_training_steps=snake_case		,	num_cycles=snake_case		,	last_epoch=snake_case		,	)

					if name == SchedulerType.POLYNOMIAL:
										return schedule_func(
										    snake_case		,	num_warmup_steps=snake_case		,	num_training_steps=snake_case		,	power=snake_case		,	last_epoch=snake_case		,	)

					return schedule_func(
					    snake_case		,	num_warmup_steps=snake_case		,	num_training_steps=snake_case		,	last_epoch=snake_case				)






