




import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
		import torch

__A    :				str    						=	logging.get_logger(__name__)
def        lowerCamelCase_     (     ):
			'''simple docstring'''
			return importlib.util.find_spec("""deepspeed""" ) is not None


if is_accelerate_available() and is_deepspeed_available():
		from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
		# Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
		# Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
		from builtins import object as DeepSpeedConfig




class       UpperCAmelCase_   ( A							):




							'''simple docstring'''



							def __init__(       self							:						str							,  a							:						Union[str, Any]      )      ->       List[str]:
										# set global weakref object
										set_hf_deepspeed_config(self      )
										dep_version_check("""accelerate"""      )
										dep_version_check("""deepspeed"""      )
										super().__init__(a      )





class       UpperCAmelCase_   ( A							):




							'''simple docstring'''



							def __init__(       self							:						Dict							,  a							:						str      )      ->       str:
										super().__init__(a      )
										SCREAMING_SNAKE_CASE     =       None
										SCREAMING_SNAKE_CASE     =       []





							def  _UpperCAmelCase  (       self							:						str      )      ->       Optional[Any]:
										if self._dtype is None:
													raise ValueError("""trainer_config_process() wasn't called yet to tell dtype"""      )
										return self._dtype





							def  _UpperCAmelCase  (       self							:						List[Any]							,  a							:						int      )      ->       Tuple:
										SCREAMING_SNAKE_CASE     =       self.get_value(a      )
										if val is None:
													return False
										else:
													return val == "auto"





							def  _UpperCAmelCase  (       self							:						Optional[Any]							,  a							:						int							,  a							:						int							,  a							:						Dict=None							,  a							:						Dict=True      )      ->       Dict:
										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       self.find_config_node(a      )
										if config is None:
													return

										if config.get(a      ) == "auto":
													SCREAMING_SNAKE_CASE     =       hf_val
													return

										if not must_match:
													return

										SCREAMING_SNAKE_CASE     =       config.get(a      )
										if ds_val is not None and ds_val != hf_val:
													self.mismatches.append(f"""- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}"""      )

							a__    	=  partialmethod(A		,   must_match=A							)





							def  _UpperCAmelCase  (       self							:						str							,  a							:						Optional[int]      )      ->       Dict:
										SCREAMING_SNAKE_CASE     =       args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
										self.fill_match(
										    """train_micro_batch_size_per_gpu"""							,  args.per_device_train_batch_size							,  """per_device_train_batch_size"""      )
										self.fill_match("""gradient_accumulation_steps"""							,  args.gradient_accumulation_steps							,  """gradient_accumulation_steps"""      )
										self.fill_match("""train_batch_size"""							,  a							,  """train_batch_size (calculated)"""      )
										self.fill_match("""gradient_clipping"""							,  args.max_grad_norm							,  """max_grad_norm"""      )

										self.fill_match("""optimizer.params.lr"""							,  args.learning_rate							,  """learning_rate"""      )
										self.fill_match("""optimizer.params.betas"""							,  [args.adam_betaa, args.adam_betaa]							,  """adam_beta1+adam_beta2"""      )
										self.fill_match("""optimizer.params.eps"""							,  args.adam_epsilon							,  """adam_epsilon"""      )
										self.fill_match("""optimizer.params.weight_decay"""							,  args.weight_decay							,  """weight_decay"""      )

										self.fill_only("""scheduler.params.warmup_min_lr"""							,  0      )  # not a trainer arg
										self.fill_match("""scheduler.params.warmup_max_lr"""							,  args.learning_rate							,  """learning_rate"""      )
										# total_num_steps - will get set in trainer_config_finalize

										# fp16
										if args.fpaa or args.fpaa_full_eval:
													SCREAMING_SNAKE_CASE     =       """apex""" if args.fpaa_backend == """apex""" else """amp"""
										else:
													SCREAMING_SNAKE_CASE     =       None

										if args.save_on_each_node:
													# deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
													SCREAMING_SNAKE_CASE     =       self.config.get("""checkpoint"""							,  {}      )
													SCREAMING_SNAKE_CASE     =       args.save_on_each_node

										# amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
										# any here unless the user did the work
										self.fill_match(
										    """fp16.enabled"""							,  ((args.fpaa or args.fpaa_full_eval) and fpaa_backend == """amp""")							,  """fp16|fp16_full_eval+fp16_backend(amp)"""							,  )

										# apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
										# ZeRO features
										self.fill_match("""amp.enabled"""							,  fpaa_backend == """apex"""							,  """fp16+fp16_backend(apex)"""      )
										self.fill_match("""amp.opt_level"""							,  args.fpaa_opt_level							,  """fp16_opt_level"""      )

										self.fill_match("""bf16.enabled"""							,  (args.bfaa or args.bfaa_full_eval)							,  """bf16|bf16_full_eval"""      )

										# deepspeed's default mode is fp16 unless there is a config that says differently
										if self.is_true("""bf16.enabled"""      ):
													SCREAMING_SNAKE_CASE     =       torch.bfloataa
										elif self.is_false("""fp16.enabled"""      ):
													SCREAMING_SNAKE_CASE     =       torch.floataa
										else:
													SCREAMING_SNAKE_CASE     =       torch.floataa







							def  _UpperCAmelCase  (       self							:						Tuple							,  a							:						Any							,  a							:						int							,  a							:						Optional[Any]      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       [
										    """zero_optimization.reduce_bucket_size""",
										    """zero_optimization.stage3_prefetch_bucket_size""",
										    """zero_optimization.stage3_param_persistence_threshold""",
										]
										SCREAMING_SNAKE_CASE     =       [x for x in hidden_size_based_keys if self.is_auto(a      )]

										if len(a      ) > 0:
													if hasattr(model.config							,  """hidden_size"""      ):
																SCREAMING_SNAKE_CASE     =       model.config.hidden_size
													elif hasattr(model.config							,  """hidden_sizes"""      ):
																# if there are many hidden sizes pick the largest one
																SCREAMING_SNAKE_CASE     =       max(model.config.hidden_sizes      )
													else:
																raise ValueError(
																    """The model's config file has neither `hidden_size` nor `hidden_sizes` entry, """
																    """therefore it's not possible to automatically fill out the following `auto` entries """
																    f"""in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing """
																    """`auto` values for these keys with an integer value of your choice."""      )

													self.fill_only("""zero_optimization.reduce_bucket_size"""							,  hidden_size * hidden_size      )
													if self.is_zeroa():
																# automatically assign the optimal config values based on model config
																self.fill_only("""zero_optimization.stage3_prefetch_bucket_size"""							,  0.9 * hidden_size * hidden_size      )
																self.fill_only("""zero_optimization.stage3_param_persistence_threshold"""							,  10 * hidden_size      )

        # scheduler
										self.fill_match("""scheduler.params.total_num_steps"""							,  a							,  """num_training_steps (calculated)"""      )
										self.fill_match("""scheduler.params.warmup_num_steps"""							,  args.get_warmup_steps(a      )							,  """warmup_steps"""      )

										if len(self.mismatches      ) > 0:
													SCREAMING_SNAKE_CASE     =       """\n""".join(self.mismatches      )
													raise ValueError(
													    """Please correct the following DeepSpeed config values that mismatch TrainingArguments"""
													    f""" values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to 'auto'."""      )





__A    :				str    						=	None
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			# this is a special weakref global object to allow us to get to Deepspeed config from APIs
			# that don't have an easy way to get to the Deepspeed config outside of the Trainer domain.
			global _hf_deepspeed_config_weak_ref
			# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
			SCREAMING_SNAKE_CASE     =       weakref.ref(SCREAMING_SNAKE_CASE )
def        lowerCamelCase_     (     ):
			'''simple docstring'''
			# useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
			global _hf_deepspeed_config_weak_ref
			SCREAMING_SNAKE_CASE     =       None
def        lowerCamelCase_     (     ):
			'''simple docstring'''
			if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
						return _hf_deepspeed_config_weak_ref().is_zeroa()
			else:
						return False
def        lowerCamelCase_     (     ):
			'''simple docstring'''
			if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
						return _hf_deepspeed_config_weak_ref().config
			else:
						return None
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			from accelerate.utils import DummyOptim, DummyScheduler

			SCREAMING_SNAKE_CASE     =       hf_deepspeed_config.config

			# Optimizer + Scheduler
			# Currently supported combos:
			# 1. DS scheduler + DS optimizer: Yes
			# 2. HF scheduler + HF optimizer: Yes
			# 3. DS scheduler + HF optimizer: Yes
			# 4. HF scheduler + DS optimizer: No
			#
			# Unless Offload is enabled in which case it's:
			# 1. DS scheduler + DS optimizer: Yes
			# 2. HF scheduler + HF optimizer: Mostly*
			# 3. DS scheduler + HF optimizer: Mostly*
			# 4. HF scheduler + DS optimizer: No
			#
			# Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

			SCREAMING_SNAKE_CASE     =       None
			if "optimizer" in config:
						if args.adafactor:
									raise ValueError(
									    """--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. """
									    """Only one optimizer can be configured.""" )
						SCREAMING_SNAKE_CASE     =       DummyOptim(params=SCREAMING_SNAKE_CASE )
			else:
						if hf_deepspeed_config.is_offload():
									logger.info(
									    """Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the"""
									    """ custom optimizer has both CPU and GPU implementation (except LAMB)""" )

						# ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
						# But trainer uses AdamW by default.
						SCREAMING_SNAKE_CASE     =       trainer.create_optimizer()
						# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
						SCREAMING_SNAKE_CASE     =       True

			SCREAMING_SNAKE_CASE     =       None
			if "scheduler" in config:
						SCREAMING_SNAKE_CASE     =       DummyScheduler(SCREAMING_SNAKE_CASE )
			else:
						if isinstance(SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
									raise ValueError(
									    """Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. """
									    """Please configure a scheduler in the DeepSpeed config.""" )
						SCREAMING_SNAKE_CASE     =       trainer.create_scheduler(num_training_steps=SCREAMING_SNAKE_CASE       ,						optimizer=SCREAMING_SNAKE_CASE )

			return optimizer, lr_scheduler
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE=False ):
			'''simple docstring'''
			from deepspeed.utils import logger as ds_logger

			SCREAMING_SNAKE_CASE     =       trainer.model
			SCREAMING_SNAKE_CASE     =       trainer.args

			SCREAMING_SNAKE_CASE     =       trainer.accelerator.state.deepspeed_plugin.hf_ds_config

			# resume config update - some bits like `model` and `num_training_steps` only become available during train
			hf_deepspeed_config.trainer_config_finalize(SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE )

			# set the Deepspeed log level consistent with the Trainer
			ds_logger.setLevel(args.get_process_log_level() )

			if inference:
						# only Z3 makes sense for the inference
						if not hf_deepspeed_config.is_zeroa():
									raise ValueError("""ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config""" )

						# in case the training config is re-used for inference
						hf_deepspeed_config.del_config_sub_tree("""optimizer""" )
						hf_deepspeed_config.del_config_sub_tree("""lr_scheduler""" )
						SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       None, None
						SCREAMING_SNAKE_CASE     =       None
			else:
						SCREAMING_SNAKE_CASE     =       None  # important for when deepspeed_init is used as re-init
						SCREAMING_SNAKE_CASE     =       list(filter(lambda SCREAMING_SNAKE_CASE : p.requires_grad       ,						model.parameters() ) )
						SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       deepspeed_optim_sched(
						    SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE )

			# keep for quick debug:
			# from pprint import pprint; pprint(config)

			return optimizer, lr_scheduler




def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			# it's possible that the user is trying to resume from model_path, which doesn't necessarily
			# contain a deepspeed checkpoint. e.g. examples just check if the dir exists and assume it's
			# a resume from a checkpoint and not just a local pretrained weight. So we check here if the
			# path contains what looks like a deepspeed checkpoint
			import glob

			SCREAMING_SNAKE_CASE     =       sorted(glob.glob(f"""{checkpoint_path}/global_step*""" ) )

			if len(SCREAMING_SNAKE_CASE ) > 0:
						logger.info(f"""Attempting to resume from {checkpoint_path}""" )
						# this magically updates self.optimizer and self.lr_scheduler
						SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       deepspeed_engine.load_checkpoint(
						    SCREAMING_SNAKE_CASE       ,						load_optimizer_states=SCREAMING_SNAKE_CASE       ,						load_lr_scheduler_states=SCREAMING_SNAKE_CASE )
						if load_path is None:
									raise ValueError(f"""[deepspeed] failed to resume from checkpoint {checkpoint_path}""" )
			else:
						raise ValueError(f"""Can't find a valid checkpoint at {checkpoint_path}""" )












import argparse

import torch

from transformers import NystromformerConfig, NystromformerForMaskedLM
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			if "model" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""model."""       ,						"""""" )
			if "norm1" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""norm1"""       ,						"""attention.output.LayerNorm""" )
			if "norm2" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""norm2"""       ,						"""output.LayerNorm""" )
			if "norm" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""norm"""       ,						"""LayerNorm""" )
			if "transformer" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.split(""".""" )[0].split("""_""" )[-1]
						SCREAMING_SNAKE_CASE     =       orig_key.replace(f"""transformer_{layer_num}"""       ,						f"""encoder.layer.{layer_num}""" )
			if "mha.attn" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""mha.attn"""       ,						"""attention.self""" )
			if "mha" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""mha"""       ,						"""attention""" )
			if "W_q" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""W_q"""       ,						"""self.query""" )
			if "W_k" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""W_k"""       ,						"""self.key""" )
			if "W_v" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""W_v"""       ,						"""self.value""" )
			if "ff1" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""ff1"""       ,						"""intermediate.dense""" )
			if "ff2" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""ff2"""       ,						"""output.dense""" )
			if "ff" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""ff"""       ,						"""output.dense""" )
			if "mlm_class" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""mlm.mlm_class"""       ,						"""cls.predictions.decoder""" )
			if "mlm" in orig_key:
						SCREAMING_SNAKE_CASE     =       orig_key.replace("""mlm"""       ,						"""cls.predictions.transform""" )
			if "cls" not in orig_key:
						SCREAMING_SNAKE_CASE     =       """nystromformer.""" + orig_key

			return orig_key
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			for key in orig_state_dict.copy().keys():
						SCREAMING_SNAKE_CASE     =       orig_state_dict.pop(SCREAMING_SNAKE_CASE )

						if ("pooler" in key) or ("sen_class" in key) or ("conv.bias" in key):
									continue
						else:
									SCREAMING_SNAKE_CASE     =       val

			SCREAMING_SNAKE_CASE     =       orig_state_dict["""cls.predictions.decoder.bias"""]
			SCREAMING_SNAKE_CASE     =       (
			    torch.arange(config.max_position_embeddings ).expand((1, -1) ) + 2
			)

			return orig_state_dict
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			SCREAMING_SNAKE_CASE     =       torch.load(SCREAMING_SNAKE_CASE       ,						map_location="""cpu""" )["""model_state_dict"""]
			SCREAMING_SNAKE_CASE     =       NystromformerConfig.from_json_file(SCREAMING_SNAKE_CASE )
			SCREAMING_SNAKE_CASE     =       NystromformerForMaskedLM(SCREAMING_SNAKE_CASE )

			SCREAMING_SNAKE_CASE     =       convert_checkpoint_helper(SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE )

			model.load_state_dict(SCREAMING_SNAKE_CASE )
			model.eval()
			model.save_pretrained(SCREAMING_SNAKE_CASE )

			print(f"""Checkpoint successfuly converted. Model saved at {pytorch_dump_path}""" )


if __name__ == "__main__":
		__A    :				List[str]    						=	argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--pytorch_model_path""", default=None, type=str, required=True, help="""Path to Nystromformer pytorch checkpoint."""
		)
		parser.add_argument(
		    """--config_file""",
		    default=None,
		    type=str,
		    required=True,
		    help="""The json file for Nystromformer model config.""",
		)
		parser.add_argument(
		    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
		)
		__A    :				Union[str, Any]    						=	parser.parse_args()
		convert_nystromformer_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)






