






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

snake_case_   :	Optional[Any]							      =logging.get_logger(__name__)



def      UpperCAmelCase						( ):


		'''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        a__      (	lowerCAmelCase__		):






		def __init__(      self							, lowercase__						)				->				Tuple:
				# set global weakref object
				set_hf_deepspeed_config(self						)
				dep_version_check("accelerate"						)
				dep_version_check("deepspeed"						)
				super().__init__(lowercase__						)





class        a__      (	lowerCAmelCase__		):






		def __init__(      self							, lowercase__						)				->				Union[str, Any]:
				super().__init__(lowercase__						)
				__A    =							None
				__A    =							[]






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






		def 					_lowerCamelCase					(      self							, lowercase__						)				->				str:
				__A    =							self.get_value(lowercase__						)
				if val is None:
						return False
				else:
						return val == "auto"






		def 					_lowerCamelCase					(      self							, lowercase__							, lowercase__							, lowercase__=None							, lowercase__=True						)				->				Dict:
				__A ,   __A    =							self.find_config_node(lowercase__						)
				if config is None:
						return

				if config.get(lowercase__						) == "auto":
						__A    =							hf_val
						return

				if not must_match:
						return

				__A    =							config.get(lowercase__						)
				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}"""						)

		UpperCAmelCase_							:       int 					= partialmethod(lowerCAmelCase__	,       must_match=lowerCAmelCase__		)






		def 					_lowerCamelCase					(      self							, lowercase__						)				->				int:
				__A    =							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"							, lowercase__							, "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:
						__A    =							"apex" if args.fpaa_backend == "apex" else "amp"
				else:
						__A    =							None

				if args.save_on_each_node:
						# deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
						__A    =							self.config.get("checkpoint"							, {}						)
						__A    =							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"						):
						__A    =							torch.bfloataa
				elif self.is_false("fp16.enabled"						):
						__A    =							torch.floataa
				else:
						__A    =							torch.floataa






		def 					_lowerCamelCase					(      self							, lowercase__							, lowercase__							, lowercase__						)				->				Dict:
				__A    =							[
				    "zero_optimization.reduce_bucket_size",
				    "zero_optimization.stage3_prefetch_bucket_size",
				    "zero_optimization.stage3_param_persistence_threshold",
				]
				__A    =							[x for x in hidden_size_based_keys if self.is_auto(lowercase__						)]

				if len(lowercase__						) > 0:
						if hasattr(model.config							, "hidden_size"						):
								__A    =							model.config.hidden_size
						elif hasattr(model.config							, "hidden_sizes"						):
								# if there are many hidden sizes pick the largest one
								__A    =							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"							, lowercase__							, "num_training_steps (calculated)"						)
				self.fill_match("scheduler.params.warmup_num_steps"							, args.get_warmup_steps(lowercase__						)							, "warmup_steps"						)

				if len(self.mismatches						) > 0:
						__A    =							"\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'."""						)



snake_case_   :	Tuple							      =None



def      UpperCAmelCase						( lowerCAmelCase__	):


		'''simple docstring'''

		global _hf_deepspeed_config_weak_ref
		# will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
		__A    =							weakref.ref(lowerCAmelCase__	)



def      UpperCAmelCase						( ):


		'''simple docstring'''

		global _hf_deepspeed_config_weak_ref
		__A    =							None



def      UpperCAmelCase						( ):


		'''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      UpperCAmelCase						( ):


		'''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      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	):


		'''simple docstring'''

		from accelerate.utils import DummyOptim, DummyScheduler

		__A    =							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)

		__A    =							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."	)
				__A    =							DummyOptim(params=lowerCAmelCase__	)
		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.
				__A    =							trainer.create_optimizer()
				# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
				__A    =							True

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

		return optimizer, lr_scheduler



def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__=False	):


		'''simple docstring'''

		from deepspeed.utils import logger as ds_logger

		__A    =							trainer.model
		__A    =							trainer.args

		__A    =							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(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	)

		# 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"	)
				__A ,   __A    =							None, None
				__A    =							None
		else:
				__A    =							None  # important for when deepspeed_init is used as re-init
				__A    =							list(filter(lambda lowerCAmelCase__	: p.requires_grad					,     model.parameters()	)	)
				__A ,   __A    =							deepspeed_optim_sched(
				    lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__	)

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

		return optimizer, lr_scheduler





def      UpperCAmelCase						( lowerCAmelCase__					,     lowerCAmelCase__	):


		'''simple docstring'''

		import glob

		__A    =							sorted(glob.glob(F"""{checkpoint_path}/global_step*"""	)	)

		if len(lowerCAmelCase__	) > 0:
				logger.info(F"""Attempting to resume from {checkpoint_path}"""	)
				# this magically updates self.optimizer and self.lr_scheduler
				__A ,   __A    =							deepspeed_engine.load_checkpoint(
				    lowerCAmelCase__					,     load_optimizer_states=lowerCAmelCase__					,     load_lr_scheduler_states=lowerCAmelCase__	)
				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}"""	)





"""simple docstring"""
import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling






def 						lowercase				(	_snake_case		: Optional[Any]	)	->Union[str, Any]:
	"""simple docstring"""





	__snake_case      :		str         =							'''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

	with open(_snake_case	,  '''wb'''	) as f:
		with requests.get(_snake_case	,  stream=_snake_case	) as req:
			for chunk in req.iter_content(chunk_size=2_048	):
				f.write(_snake_case	)






def 						lowercase				(	)	->VivitConfig:
	"""simple docstring"""





	__snake_case      :		int         =							VivitConfig()

	__snake_case      :		Tuple         =							400
	__snake_case      :		Tuple         =							'''huggingface/label-files'''
	__snake_case      :		List[str]         =							'''kinetics400-id2label.json'''

	__snake_case      :		Dict         =							json.load(open(hf_hub_download(_snake_case	,  _snake_case	,  repo_type='''dataset'''	)	,  '''r'''	)	)
	__snake_case      :		List[str]         =							{int(_snake_case	): v for k, v in idalabel.items()}
	__snake_case      :		Tuple         =							idalabel
	__snake_case      :		Optional[int]         =							{v: k for k, v in idalabel.items()}
	return config






def 						lowercase				(	)	->Optional[int]:
	"""simple docstring"""





	__snake_case      :		List[Any]         =							hf_hub_download(
	    repo_id='''hf-internal-testing/spaghetti-video'''	,  filename='''eating_spaghetti_32_frames.npy'''	,  repo_type='''dataset'''	)
	__snake_case      :		str         =							np.load(_snake_case	)
	return list(_snake_case	)






def 						lowercase				(	_snake_case		: np.ndarray	)	->str:
	"""simple docstring"""





	if np.ndim(_snake_case	) == 2:
		return transform_attention_bias(_snake_case	)

	elif np.ndim(_snake_case	) == 3:
		return transform_attention_kernel(_snake_case	)

	else:
		raise Exception(f"""Invalid number of dimesions: {np.ndim(_snake_case	)}"""	)






def 						lowercase				(	_snake_case		: np.ndarray	)	->Tuple:
	"""simple docstring"""





	return current.flatten()






def 						lowercase				(	_snake_case		: np.ndarray	)	->List[str]:
	"""simple docstring"""





	return np.reshape(_snake_case	,  (current.shape[0], current.shape[1] * current.shape[2])	).T






def 						lowercase				(	_snake_case		: np.ndarray	)	->Optional[Any]:
	"""simple docstring"""





	return np.reshape(_snake_case	,  (current.shape[0] * current.shape[1], current.shape[2])	).T






def 						lowercase				(	_snake_case		: Dict	,  _snake_case		: Tuple	)	->Tuple:
	"""simple docstring"""





	__snake_case      :		int         =							state_dict['''optimizer''']['''target''']['''Transformer'''][f"""encoderblock_{i}"""]

	__snake_case      :		Optional[int]         =							f"""encoder.layer.{i}."""
	__snake_case      :		Dict         =							{
	    prefix + '''intermediate.dense.bias''': state['''MlpBlock_0''']['''Dense_0''']['''bias'''],
	    prefix + '''intermediate.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_0''']['''kernel''']	),
	    prefix + '''output.dense.bias''': state['''MlpBlock_0''']['''Dense_1''']['''bias'''],
	    prefix + '''output.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_1''']['''kernel''']	),
	    prefix + '''layernorm_before.bias''': state['''LayerNorm_0''']['''bias'''],
	    prefix + '''layernorm_before.weight''': state['''LayerNorm_0''']['''scale'''],
	    prefix + '''layernorm_after.bias''': state['''LayerNorm_1''']['''bias'''],
	    prefix + '''layernorm_after.weight''': state['''LayerNorm_1''']['''scale'''],
	    prefix
	    + '''attention.attention.query.bias''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''query''']['''bias''']	),
	    prefix
	    + '''attention.attention.query.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''query''']['''kernel''']	),
	    prefix
	    + '''attention.attention.key.bias''': transform_attention(state['''MultiHeadDotProductAttention_0''']['''key''']['''bias''']	),
	    prefix
	    + '''attention.attention.key.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''key''']['''kernel''']	),
	    prefix
	    + '''attention.attention.value.bias''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''value''']['''bias''']	),
	    prefix
	    + '''attention.attention.value.weight''': transform_attention(
	        state['''MultiHeadDotProductAttention_0''']['''value''']['''kernel''']	),
	    prefix + '''attention.output.dense.bias''': state['''MultiHeadDotProductAttention_0''']['''out''']['''bias'''],
	    prefix
	    + '''attention.output.dense.weight''': transform_attention_output_weight(
	        state['''MultiHeadDotProductAttention_0''']['''out''']['''kernel''']	),
	}

	return new_state






def 						lowercase				(	_snake_case		: int	)	->List[str]:
	"""simple docstring"""





	return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]	)






def 						lowercase				(	_snake_case		: Optional[Any]	,  _snake_case		: Optional[int]=False	)	->Union[str, Any]:
	"""simple docstring"""





	__snake_case      :		Tuple         =							get_n_layers(_snake_case	)

	__snake_case      :		str         =							OrderedDict()

	__snake_case      :		List[Any]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
	__snake_case      :		List[str]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']

	__snake_case      :		Optional[Any]         =							np.transpose(
	    state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']	,  (4, 3, 0, 1, 2)	)
	__snake_case      :		Optional[Any]         =							state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']

	__snake_case      :		Union[str, Any]         =							state_dict['''optimizer''']['''target''']['''cls''']
	__snake_case      :		List[str]         =							state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
	    '''pos_embedding'''
	]

	for i in range(_snake_case	):
		new_state.update(transform_state_encoder_block(_snake_case	,  _snake_case	)	)

	if classification_head:
		__snake_case      :		List[Any]         =							{'''vivit.''' + k: v for k, v in new_state.items()}
		__snake_case      :		str         =							np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']	)
		__snake_case      :		int         =							np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''bias''']	)

	return {k: torch.tensor(_snake_case	) for k, v in new_state.items()}






def 						lowercase				(	)	->VivitImageProcessor:
	"""simple docstring"""





	__snake_case      :		List[str]         =							VivitImageProcessor()

	assert extractor.do_resize is True
	assert extractor.size == {"shortest_edge": 256}
	assert extractor.do_center_crop is True
	assert extractor.crop_size == {"width": 224, "height": 224}
	assert extractor.resample == PILImageResampling.BILINEAR

	# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
	# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
	# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
	assert extractor.do_normalize is False
	assert extractor.do_rescale is True
	assert extractor.rescale_factor == 1 / 255

	# zero-centering = True in original implementation
	assert extractor.do_zero_centering is True

	return extractor






def 						lowercase				(	_snake_case		: str	)	->Optional[Any]:
	"""simple docstring"""





	__snake_case      :		Tuple         =							'''checkpoint'''

	if not os.path.exists(_snake_case	):
		download_checkpoint(_snake_case	)

	__snake_case      :		List[Any]         =							restore_checkpoint(_snake_case	,  _snake_case	)
	__snake_case      :		Union[str, Any]         =							transform_state(_snake_case	,  classification_head=_snake_case	)

	__snake_case      :		Optional[Any]         =							get_vivit_config()

	assert config.image_size == 224
	assert config.num_frames == 32

	__snake_case      :		List[str]         =							VivitForVideoClassification(_snake_case	)
	model.load_state_dict(_snake_case	)
	model.eval()

	__snake_case      :		Tuple         =							get_processor()

	__snake_case      :		Any         =							prepare_video()
	__snake_case      :		List[str]         =							extractor(_snake_case	,  return_tensors='''pt'''	)

	__snake_case      :		Union[str, Any]         =							model(**_snake_case	)

	__snake_case      :		List[str]         =							torch.Size([1, 400]	)
	__snake_case      :		Optional[int]         =							torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]	)

	assert outputs.logits.shape == expected_shape
	assert torch.allclose(outputs.logits[0, :5]	,  _snake_case	,  atol=1e-4	), outputs.logits[0, :5]

	model.save_pretrained(_snake_case	)
	extractor.save_pretrained(_snake_case	)


if __name__ == "__main__":
			SCREAMING_SNAKE_CASE					:       Optional[Any]							   =							argparse.ArgumentParser()

			parser.add_argument("""--output_model_name""", """-o""", type=str, help="""Output path for the converted HuggingFace model""")

			SCREAMING_SNAKE_CASE					:       List[str]							   =							parser.parse_args()
			convert(args.output_model_name)





