






'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase_				: Dict    	=							{'''configuration_xlnet''': ['''XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''XLNetConfig''']}

try:
					if not is_sentencepiece_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase_				: Dict    	=							['''XLNetTokenizer''']

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase_				: str    	=							['''XLNetTokenizerFast''']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase_				: Union[str, Any]    	=							[
					    '''XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''XLNetForMultipleChoice''',
					    '''XLNetForQuestionAnswering''',
					    '''XLNetForQuestionAnsweringSimple''',
					    '''XLNetForSequenceClassification''',
					    '''XLNetForTokenClassification''',
					    '''XLNetLMHeadModel''',
					    '''XLNetModel''',
					    '''XLNetPreTrainedModel''',
					    '''load_tf_weights_in_xlnet''',
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase_				: List[Any]    	=							[
					    '''TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''TFXLNetForMultipleChoice''',
					    '''TFXLNetForQuestionAnsweringSimple''',
					    '''TFXLNetForSequenceClassification''',
					    '''TFXLNetForTokenClassification''',
					    '''TFXLNetLMHeadModel''',
					    '''TFXLNetMainLayer''',
					    '''TFXLNetModel''',
					    '''TFXLNetPreTrainedModel''',
					]


if TYPE_CHECKING:
					from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig

					try:
										if not is_sentencepiece_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_xlnet import XLNetTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_xlnet_fast import XLNetTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_xlnet import (
										    XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
										    XLNetForMultipleChoice,
										    XLNetForQuestionAnswering,
										    XLNetForQuestionAnsweringSimple,
										    XLNetForSequenceClassification,
										    XLNetForTokenClassification,
										    XLNetLMHeadModel,
										    XLNetModel,
										    XLNetPreTrainedModel,
										    load_tf_weights_in_xlnet,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_xlnet import (
										    TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFXLNetForMultipleChoice,
										    TFXLNetForQuestionAnsweringSimple,
										    TFXLNetForSequenceClassification,
										    TFXLNetForTokenClassification,
										    TFXLNetLMHeadModel,
										    TFXLNetMainLayer,
										    TFXLNetModel,
										    TFXLNetPreTrainedModel,
										)

else:
					import sys

					lowerCamelCase_				: Optional[Any]    	=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)








'''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    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						Optional[int] ,      lowercase							:						Optional[int] ,      lowercase							:						Optional[Any] ,      lowercase							:						bool = True ,      lowercase							:						bool = False			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = scheduler
									UpperCamelCase__         = optimizers if isinstance(lowercase ,      (list, tuple)			) else [optimizers]
									UpperCamelCase__         = split_batches
									UpperCamelCase__         = step_with_optimizer
									UpperCamelCase__         = GradientState()






			def 				A      (   self							:						int ,      *lowercase							:						int ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									if not self.step_with_optimizer:
															# No link between scheduler and optimizer -> just step
															self.scheduler.step(*lowercase ,      **lowercase			)
															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(*lowercase ,      **lowercase			)
									else:
															# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
															# num_processes steps per training step
															UpperCamelCase__         = AcceleratorState().num_processes
															for _ in range(lowercase			):
																					# 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(*lowercase ,      **lowercase			)
																					else:
																											self.scheduler.step(*lowercase ,      **lowercase			)






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

									'''simple docstring'''







									return self.scheduler.get_last_lr()






			def 				A      (   self							:						Any			)							->    int:

									'''simple docstring'''







									return self.scheduler.state_dict()






			def 				A      (   self							:						Any ,      lowercase							:						int			)							->    Optional[int]:

									'''simple docstring'''







									self.scheduler.load_state_dict(lowercase			)






			def 				A      (   self							:						str			)							->    int:

									'''simple docstring'''







									return self.scheduler.get_lr()







			def 				A      (   self							:						Union[str, Any] ,      *lowercase							:						List[Any] ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									return self.scheduler.print_lr(*lowercase ,      **lowercase			)
