



"""simple docstring"""






import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('>=', FSDP_PYTORCH_VERSION):
	import torch.distributed.checkpoint as dist_cp
	from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
	from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
	from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
	from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


__snake_case				:  Dict =       get_logger(__name__)


def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case=0      )		->  Any:
			os.makedirs(__snake_case   ,exist_ok=__snake_case      )
			with FSDP.state_dict_type(
			    __snake_case   ,fsdp_plugin.state_dict_type   ,fsdp_plugin.state_dict_config   ,fsdp_plugin.optim_state_dict_config      ):
						__lowerCAmelCase :    Tuple				=			model.state_dict()
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									__lowerCAmelCase :    Optional[int]				=			F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
									__lowerCAmelCase :    List[Any]				=			os.path.join(__snake_case   ,__snake_case      )
									if accelerator.process_index == 0:
												logger.info(F"""Saving model to {output_model_file}"""      )
												torch.save(__snake_case   ,__snake_case      )
												logger.info(F"""Model saved to {output_model_file}"""      )
						elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
									__lowerCAmelCase :    List[Any]				=			(
									    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
									    if model_index == 0
									    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
									)
									__lowerCAmelCase :    Any				=			os.path.join(__snake_case   ,__snake_case      )
									logger.info(F"""Saving model to {output_model_file}"""      )
									torch.save(__snake_case   ,__snake_case      )
									logger.info(F"""Model saved to {output_model_file}"""      )
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									__lowerCAmelCase :    str				=			os.path.join(__snake_case   ,F"""{MODEL_NAME}_{model_index}"""      )
									os.makedirs(__snake_case   ,exist_ok=__snake_case      )
									logger.info(F"""Saving model to {ckpt_dir}"""      )
									__lowerCAmelCase :    Union[str, Any]				=			{"model": state_dict}

									dist_cp.save_state_dict(
									    state_dict=__snake_case   ,storage_writer=dist_cp.FileSystemWriter(__snake_case      )   ,planner=DefaultSavePlanner()   ,)
									logger.info(F"""Model saved to {ckpt_dir}"""      )


def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case=0      )		->  Union[str, Any]:
			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    __snake_case   ,fsdp_plugin.state_dict_type   ,fsdp_plugin.state_dict_config   ,fsdp_plugin.optim_state_dict_config      ):
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									if type(__snake_case      ) != FSDP and accelerator.process_index != 0:
												if not fsdp_plugin.sync_module_states:
															raise ValueError(
															    "Set the `sync_module_states` flag to `True` so that model states are synced across processes when "
															    "initializing FSDP object"      )
												return
									__lowerCAmelCase :    Optional[int]				=			F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
									__lowerCAmelCase :    Optional[int]				=			os.path.join(__snake_case   ,__snake_case      )
									logger.info(F"""Loading model from {input_model_file}"""      )
									__lowerCAmelCase :    Dict				=			torch.load(__snake_case      )
									logger.info(F"""Model loaded from {input_model_file}"""      )
						elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
									__lowerCAmelCase :    str				=			(
									    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
									    if model_index == 0
									    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
									)
									__lowerCAmelCase :    str				=			os.path.join(__snake_case   ,__snake_case      )
									logger.info(F"""Loading model from {input_model_file}"""      )
									__lowerCAmelCase :    List[Any]				=			torch.load(__snake_case      )
									logger.info(F"""Model loaded from {input_model_file}"""      )
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									__lowerCAmelCase :    int				=			(
									    os.path.join(__snake_case   ,F"""{MODEL_NAME}_{model_index}"""      )
									    if F"""{MODEL_NAME}""" not in input_dir
									    else input_dir
									)
									logger.info(F"""Loading model from {ckpt_dir}"""      )
									__lowerCAmelCase :    List[str]				=			{"model": model.state_dict()}
									dist_cp.load_state_dict(
									    state_dict=__snake_case   ,storage_reader=dist_cp.FileSystemReader(__snake_case      )   ,planner=DefaultLoadPlanner()   ,)
									__lowerCAmelCase :    Dict				=			state_dict["model"]
									logger.info(F"""Model loaded from {ckpt_dir}"""      )
						model.load_state_dict(__snake_case      )


def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case=0      )		->  Union[str, Any]:
			os.makedirs(__snake_case   ,exist_ok=__snake_case      )
			with FSDP.state_dict_type(
			    __snake_case   ,fsdp_plugin.state_dict_type   ,fsdp_plugin.state_dict_config   ,fsdp_plugin.optim_state_dict_config      ):
						__lowerCAmelCase :    int				=			FSDP.optim_state_dict(__snake_case   ,__snake_case      )
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									if accelerator.process_index == 0:
												__lowerCAmelCase :    Any				=			(
												    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
												)
												__lowerCAmelCase :    List[Any]				=			os.path.join(__snake_case   ,__snake_case      )
												logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""      )
												torch.save(__snake_case   ,__snake_case      )
												logger.info(F"""Optimizer state saved in {output_optimizer_file}"""      )
						else:
									__lowerCAmelCase :    Any				=			os.path.join(__snake_case   ,F"""{OPTIMIZER_NAME}_{optimizer_index}"""      )
									os.makedirs(__snake_case   ,exist_ok=__snake_case      )
									logger.info(F"""Saving Optimizer state to {ckpt_dir}"""      )
									dist_cp.save_state_dict(
									    state_dict={"optimizer": optim_state}   ,storage_writer=dist_cp.FileSystemWriter(__snake_case      )   ,planner=DefaultSavePlanner()   ,)
									logger.info(F"""Optimizer state saved in {ckpt_dir}"""      )





def      _lowercase							(    __snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case   ,__snake_case=0      )		->  List[str]:
			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    __snake_case   ,fsdp_plugin.state_dict_type   ,fsdp_plugin.state_dict_config   ,fsdp_plugin.optim_state_dict_config      ):
						if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
									__lowerCAmelCase :    Union[str, Any]				=			None
									# below check should work but currently it isn't working (mostly opytorch issue),
									# in the meantime disabling it at the cost of excess memory usage
									# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
									__lowerCAmelCase :    Tuple				=			(
									    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
									)
									__lowerCAmelCase :    Union[str, Any]				=			os.path.join(__snake_case   ,__snake_case      )
									logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""      )
									__lowerCAmelCase :    Any				=			torch.load(__snake_case      )
									logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""      )
						else:
									__lowerCAmelCase :    str				=			(
									    os.path.join(__snake_case   ,F"""{OPTIMIZER_NAME}_{optimizer_index}"""      )
									    if F"""{OPTIMIZER_NAME}""" not in input_dir
									    else input_dir
									)
									logger.info(F"""Loading Optimizer from {ckpt_dir}"""      )
									__lowerCAmelCase :    Dict				=			load_sharded_optimizer_state_dict(
									    model_state_dict=model.state_dict()   ,optimizer_key="optimizer"   ,storage_reader=dist_cp.FileSystemReader(__snake_case      )   ,)
									__lowerCAmelCase :    List[str]				=			optim_state["optimizer"]
									logger.info(F"""Optimizer loaded from {ckpt_dir}"""      )
						__lowerCAmelCase :    List[str]				=			FSDP.optim_state_dict_to_load(__snake_case   ,__snake_case   ,__snake_case      )
						optimizer.load_state_dict(__snake_case      )




"""simple docstring"""






import warnings

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


__snake_case				:  Tuple =       logging.get_logger(__name__)

__snake_case				:  Tuple =       {
    '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      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      'xlnet'
	SCREAMING_SNAKE_CASE        =      ['mems']
	SCREAMING_SNAKE_CASE        =      {
	    'n_token': 'vocab_size',  # Backward compatibility
	    'hidden_size': 'd_model',
	    'num_attention_heads': 'n_head',
	    'num_hidden_layers': 'n_layer',
	}

	def __init__(	self:						Tuple			,     _SCREAMING_SNAKE_CASE:						List[str]=3_2000			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=1024			,     _SCREAMING_SNAKE_CASE:						Dict=24			,     _SCREAMING_SNAKE_CASE:						Optional[int]=16			,     _SCREAMING_SNAKE_CASE:						List[str]=4096			,     _SCREAMING_SNAKE_CASE:						Tuple="gelu"			,     _SCREAMING_SNAKE_CASE:						int=True			,     _SCREAMING_SNAKE_CASE:						List[Any]="bi"			,     _SCREAMING_SNAKE_CASE:						Any=0.02			,     _SCREAMING_SNAKE_CASE:						List[str]=1e-12			,     _SCREAMING_SNAKE_CASE:						Any=0.1			,     _SCREAMING_SNAKE_CASE:						Any=512			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=None			,     _SCREAMING_SNAKE_CASE:						Any=True			,     _SCREAMING_SNAKE_CASE:						Any=False			,     _SCREAMING_SNAKE_CASE:						Tuple=False			,     _SCREAMING_SNAKE_CASE:						List[Any]=-1			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=False			,     _SCREAMING_SNAKE_CASE:						Optional[int]="last"			,     _SCREAMING_SNAKE_CASE:						Tuple=True			,     _SCREAMING_SNAKE_CASE:						Dict="tanh"			,     _SCREAMING_SNAKE_CASE:						int=0.1			,     _SCREAMING_SNAKE_CASE:						Tuple=5			,     _SCREAMING_SNAKE_CASE:						List[str]=5			,     _SCREAMING_SNAKE_CASE:						Optional[Any]=5			,     _SCREAMING_SNAKE_CASE:						Optional[int]=1			,     _SCREAMING_SNAKE_CASE:						Tuple=2			,     **_SCREAMING_SNAKE_CASE:						str			,     )							->			Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			vocab_size
				__lowerCAmelCase :    Optional[int]				=			d_model
				__lowerCAmelCase :    Union[str, Any]				=			n_layer
				__lowerCAmelCase :    List[Any]				=			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})""")
				__lowerCAmelCase :    Optional[Any]				=			d_model // n_head
				__lowerCAmelCase :    Tuple				=			ff_activation
				__lowerCAmelCase :    Union[str, Any]				=			d_inner
				__lowerCAmelCase :    Any				=			untie_r
				__lowerCAmelCase :    int				=			attn_type

				__lowerCAmelCase :    Any				=			initializer_range
				__lowerCAmelCase :    List[Any]				=			layer_norm_eps

				__lowerCAmelCase :    int				=			dropout
				__lowerCAmelCase :    List[Any]				=			mem_len
				__lowerCAmelCase :    str				=			reuse_len
				__lowerCAmelCase :    Any				=			bi_data
				__lowerCAmelCase :    Optional[int]				=			clamp_len
				__lowerCAmelCase :    Union[str, Any]				=			same_length

				__lowerCAmelCase :    List[str]				=			summary_type
				__lowerCAmelCase :    Optional[int]				=			summary_use_proj
				__lowerCAmelCase :    Tuple				=			summary_activation
				__lowerCAmelCase :    Dict				=			summary_last_dropout
				__lowerCAmelCase :    str				=			start_n_top
				__lowerCAmelCase :    List[str]				=			end_n_top

				__lowerCAmelCase :    str				=			bos_token_id
				__lowerCAmelCase :    Optional[Any]				=			pad_token_id
				__lowerCAmelCase :    Optional[Any]				=			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."			,     _SCREAMING_SNAKE_CASE			,     )
							__lowerCAmelCase :    Optional[int]				=			kwargs["use_cache"]

				__lowerCAmelCase :    Optional[int]				=			use_mems_eval
				__lowerCAmelCase :    Any				=			use_mems_train
				super().__init__(pad_token_id=_SCREAMING_SNAKE_CASE			,     bos_token_id=_SCREAMING_SNAKE_CASE			,     eos_token_id=_SCREAMING_SNAKE_CASE			,     **_SCREAMING_SNAKE_CASE)

	@property
	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			int:

				"""simple docstring"""

				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 				_SCREAMING_SNAKE_CASE  (	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						Optional[int])							->			Optional[int]:

				"""simple docstring"""

				raise NotImplementedError(
				    F"""The model {self.model_type} is one of the few models that has no sequence length limit.""")