




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


UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			{
    """alibaba-damo/mgp-str-base""": """https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json""",
}
class 			a__    (					snake_case__      ):
						_a							:    Union[str, Any]					=    """mgp-str"""






						def __init__(   self		,			_A=[3_2, 1_2_8]		,			_A=4		,			_A=3		,			_A=2_7		,			_A=3_8		,			_A=5_0_2_5_7		,			_A=3_0_5_2_2		,			_A=7_6_8		,			_A=1_2		,			_A=1_2		,			_A=4.0		,			_A=True		,			_A=False		,			_A=1E-5		,			_A=0.0		,			_A=0.0		,			_A=0.0		,			_A=False		,			_A=0.02		,			**_A		,			):




											"""simple docstring"""



											super().__init__(**_A			)

											__lowerCAmelCase				   =       image_size
											__lowerCAmelCase				   =       patch_size
											__lowerCAmelCase				   =       num_channels
											__lowerCAmelCase				   =       max_token_length
											__lowerCAmelCase				   =       num_character_labels
											__lowerCAmelCase				   =       num_bpe_labels
											__lowerCAmelCase				   =       num_wordpiece_labels
											__lowerCAmelCase				   =       hidden_size
											__lowerCAmelCase				   =       num_hidden_layers
											__lowerCAmelCase				   =       num_attention_heads
											__lowerCAmelCase				   =       mlp_ratio
											__lowerCAmelCase				   =       distilled
											__lowerCAmelCase				   =       layer_norm_eps
											__lowerCAmelCase				   =       drop_rate
											__lowerCAmelCase				   =       qkv_bias
											__lowerCAmelCase				   =       attn_drop_rate
											__lowerCAmelCase				   =       drop_path_rate
											__lowerCAmelCase				   =       output_aa_attentions
											__lowerCAmelCase				   =       initializer_range










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


UpperCamelCase__    			=			get_logger(__name__)







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

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







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Any=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_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(SCREAMING_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				   =       F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model loaded from {input_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
															__lowerCAmelCase				   =       (
															    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
															    if model_index == 0
															    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
															)
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model loaded from {input_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_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				   =       {"model": model.state_dict()}
															dist_cp.load_state_dict(
															    state_dict=SCREAMING_SNAKE_CASE_      ,      storage_reader=dist_cp.FileSystemReader(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultLoadPlanner()      ,      )
															__lowerCAmelCase				   =       state_dict["model"]
															logger.info(F"""Model loaded from {ckpt_dir}"""     )
										model.load_state_dict(SCREAMING_SNAKE_CASE_     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Optional[Any]      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				str=0     ):
					os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										__lowerCAmelCase				   =       FSDP.optim_state_dict(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
															if accelerator.process_index == 0:
																				__lowerCAmelCase				   =       (
																				    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
																				)
																				__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""     )
																				torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Optimizer state saved in {output_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{OPTIMIZER_NAME}_{optimizer_index}"""     )
															os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_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(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultSavePlanner()      ,      )
															logger.info(F"""Optimizer state saved in {ckpt_dir}"""     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				Dict=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_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				   =       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				   =       (
															    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
															)
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_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				   =       load_sharded_optimizer_state_dict(
															    model_state_dict=model.state_dict()      ,      optimizer_key="optimizer"      ,      storage_reader=dist_cp.FileSystemReader(SCREAMING_SNAKE_CASE_     )      ,      )
															__lowerCAmelCase				   =       optim_state["optimizer"]
															logger.info(F"""Optimizer loaded from {ckpt_dir}"""     )
										__lowerCAmelCase				   =       FSDP.optim_state_dict_to_load(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										optimizer.load_state_dict(SCREAMING_SNAKE_CASE_     )




