




'''simple docstring'''





# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('''ignore''', category=UserWarning, module='''torch.optim.lr_scheduler''')

class   _a							:






     def __init__(					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE = True   ,_SCREAMING_SNAKE_CASE = False       )				->     Union[str, Any]:
          _snake_case				= scheduler
          _snake_case				= optimizers if isinstance(_SCREAMING_SNAKE_CASE   ,(list, tuple)       ) else [optimizers]
          _snake_case				= split_batches
          _snake_case				= step_with_optimizer
          _snake_case				= GradientState()






     def 							_lowercase      (					self   ,*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )				->     Union[str, Any]:
          if not self.step_with_optimizer:
               # No link between scheduler and optimizer -> just step
               self.scheduler.step(*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )
               return

          # Otherwise, first make sure the optimizer was stepped.
          if not self.gradient_state.sync_gradients:
               if self.gradient_state.adjust_scheduler:
                    self.scheduler._step_count += 1
               return

          for opt in self.optimizers:
               if opt.step_was_skipped:
                    return
          if self.split_batches:
               # Split batches -> the training dataloader batch size is not changed so one step per training step
               self.scheduler.step(*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )
          else:
               # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
               # num_processes steps per training step
               _snake_case				= AcceleratorState().num_processes
               for _ in range(_SCREAMING_SNAKE_CASE       ):
                    # Special case when using OneCycle and `drop_last` was not used
                    if hasattr(self.scheduler   ,"total_steps"       ):
                         if self.scheduler._step_count <= self.scheduler.total_steps:
                              self.scheduler.step(*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )
                    else:
                         self.scheduler.step(*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     List[Any]:
          return self.scheduler.get_last_lr()






     def 							_lowercase      (					self       )				->     int:
          return self.scheduler.state_dict()






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE       )				->     Dict:
          self.scheduler.load_state_dict(_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     Any:
          return self.scheduler.get_lr()






     def 							_lowercase      (					self   ,*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )				->     str:
          return self.scheduler.print_lr(*_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )








'''simple docstring'''





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


UpperCamelCase_		:     List[Any]          =		logging.get_logger(__name__)

UpperCamelCase_		:     Tuple          =		{
    '''RWKV/rwkv-4-169m-pile''': '''https://huggingface.co/RWKV/rwkv-4-169m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-430m-pile''': '''https://huggingface.co/RWKV/rwkv-4-430m-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-1b5-pile''': '''https://huggingface.co/RWKV/rwkv-4-1b5-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-3b-pile''': '''https://huggingface.co/RWKV/rwkv-4-3b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-7b-pile''': '''https://huggingface.co/RWKV/rwkv-4-7b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-4-14b-pile''': '''https://huggingface.co/RWKV/rwkv-4-14b-pile/resolve/main/config.json''',
    '''RWKV/rwkv-raven-1b5''': '''https://huggingface.co/RWKV/rwkv-raven-1b5/resolve/main/config.json''',
    '''RWKV/rwkv-raven-3b''': '''https://huggingface.co/RWKV/rwkv-raven-3b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-7b''': '''https://huggingface.co/RWKV/rwkv-raven-7b/resolve/main/config.json''',
    '''RWKV/rwkv-raven-14b''': '''https://huggingface.co/RWKV/rwkv-raven-14b/resolve/main/config.json''',
}

class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		"""rwkv"""
     SCREAMING_SNAKE_CASE_				:       Any       							=		{"""max_position_embeddings""": """context_length"""}






     def __init__(					self   ,_SCREAMING_SNAKE_CASE=50_277   ,_SCREAMING_SNAKE_CASE=1_024   ,_SCREAMING_SNAKE_CASE=4_096   ,_SCREAMING_SNAKE_CASE=32   ,_SCREAMING_SNAKE_CASE=None   ,_SCREAMING_SNAKE_CASE=None   ,_SCREAMING_SNAKE_CASE=1e-5   ,_SCREAMING_SNAKE_CASE=0   ,_SCREAMING_SNAKE_CASE=0   ,_SCREAMING_SNAKE_CASE=6   ,_SCREAMING_SNAKE_CASE=False   ,_SCREAMING_SNAKE_CASE=True   ,**_SCREAMING_SNAKE_CASE   ,)				->     List[str]:
          _snake_case				= vocab_size
          _snake_case				= context_length
          _snake_case				= hidden_size
          _snake_case				= num_hidden_layers
          _snake_case				= attention_hidden_size if attention_hidden_size is not None else hidden_size
          _snake_case				= intermediate_size if intermediate_size is not None else 4 * hidden_size
          _snake_case				= layer_norm_epsilon
          _snake_case				= rescale_every
          _snake_case				= use_cache

          _snake_case				= bos_token_id
          _snake_case				= eos_token_id

          super().__init__(
              tie_word_embeddings=_SCREAMING_SNAKE_CASE   ,bos_token_id=_SCREAMING_SNAKE_CASE   ,eos_token_id=_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )


