



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

__UpperCAmelCase    	=  logging.get_logger(__name__)





def     __lowerCamelCase							(   ):
							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 						lowerCamelCase__			(		_a ):


						def __init__(			self    :	str   ,							_a    :	Dict      ):
													# set global weakref object
													set_hf_deepspeed_config(self      )
													dep_version_check("accelerate"      )
													dep_version_check("deepspeed"      )
													super().__init__(_a      )


class 						lowerCamelCase__			(		_a ):


						def __init__(			self    :	List[str]   ,							_a    :	List[str]      ):
													super().__init__(_a      )
													a__:							List[str]									=None
													a__:							Any									=[]


						def 	_lowerCamelCase      (			self    :	str      ):
													if self._dtype is None:
																				raise ValueError("trainer_config_process() wasn't called yet to tell dtype"      )
													return self._dtype


						def 	_lowerCamelCase      (			self    :	Union[str, Any]   ,							_a    :	List[str]      ):
													a__:							Dict									=self.get_value(_a      )
													if val is None:
																				return False
													else:
																				return val == "auto"


						def 	_lowerCamelCase      (			self    :	Any   ,							_a    :	List[str]   ,							_a    :	List[str]   ,							_a    :	List[Any]=None   ,							_a    :	Tuple=True      ):
													a__				,      a__:							Dict									=self.find_config_node(_a      )
													if config is None:
																				return

													if config.get(_a      ) == "auto":
																				a__:							Optional[int]									=hf_val
																				return

													if not must_match:
																				return

													a__:							Any									=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}"      )

						_lowerCAmelCase             =			partialmethod(_a					,       must_match=_a )


						def 	_lowerCamelCase      (			self    :	List[str]   ,							_a    :	Any      ):
													a__:							Tuple									=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:
																				a__:							List[Any]									="apex" if args.fpaa_backend == "apex" else "amp"
													else:
																				a__:							Any									=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__:							Dict									=self.config.get("checkpoint"   ,							{}      )
																				a__:							Optional[Any]									=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__:							int									=torch.bfloataa
													elif self.is_false("fp16.enabled"      ):
																				a__:							Union[str, Any]									=torch.floataa
													else:
																				a__:							List[Any]									=torch.floataa







						def 	_lowerCamelCase      (			self    :	Tuple   ,							_a    :	Optional[Any]   ,							_a    :	Optional[int]   ,							_a    :	List[str]      ):
													a__:							str									=[
													    "zero_optimization.reduce_bucket_size",
													    "zero_optimization.stage3_prefetch_bucket_size",
													    "zero_optimization.stage3_param_persistence_threshold",
													]
													a__:							List[Any]									=[x for x in hidden_size_based_keys if self.is_auto(_a      )]

													if len(_a      ) > 0:
																				if hasattr(model.config   ,							"hidden_size"      ):
																											a__:							str									=model.config.hidden_size
																				elif hasattr(model.config   ,							"hidden_sizes"      ):
																											# if there are many hidden sizes pick the largest one
																											a__:							Union[str, Any]									=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"   ,							1_0 * 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:
																				a__:							Tuple									="\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'."      )


__UpperCAmelCase    	=  None





def     __lowerCamelCase							(   __magic_name__						:	Tuple    ):
							# 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)
							a__:							List[str]									=weakref.ref(__magic_name__    )





def     __lowerCamelCase							(   ):
							# useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
							global _hf_deepspeed_config_weak_ref
							a__:							str									=None





def     __lowerCamelCase							(   ):
							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							(   ):
							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							(   __magic_name__						:	List[str]  ,			__magic_name__						:	Any  ,			__magic_name__						:	Dict  ,			__magic_name__						:	Dict  ,			__magic_name__						:	str    ):
							from accelerate.utils import DummyOptim, DummyScheduler

							a__:							Optional[int]									=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__:							Optional[Any]									=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__:							Union[str, Any]									=DummyOptim(params=__magic_name__    )
							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__:							Tuple									=trainer.create_optimizer()
														# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
														a__:							Tuple									=True

							a__:							Optional[Any]									=None
							if "scheduler" in config:
														a__:							str									=DummyScheduler(__magic_name__    )
							else:
														if isinstance(__magic_name__  ,			__magic_name__    ):
																					raise ValueError(
																					    "Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. "
																					    "Please configure a scheduler in the DeepSpeed config."    )
														a__:							Optional[int]									=trainer.create_scheduler(num_training_steps=__magic_name__  ,			optimizer=__magic_name__    )

							return optimizer, lr_scheduler





def     __lowerCamelCase							(   __magic_name__						:	Dict  ,			__magic_name__						:	List[str]  ,			__magic_name__						:	Dict=False    ):
							from deepspeed.utils import logger as ds_logger

							a__:							str									=trainer.model
							a__:							str									=trainer.args

							a__:							Optional[Any]									=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(__magic_name__  ,			__magic_name__  ,			__magic_name__    )

							# 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__:							Any									=None, None
														a__:							Dict									=None
							else:
														a__:							List[Any]									=None  # important for when deepspeed_init is used as re-init
														a__:							Dict									=list(filter(lambda __magic_name__    : p.requires_grad  ,			model.parameters()    )    )
														a__				,      a__:							int									=deepspeed_optim_sched(
														    __magic_name__  ,			__magic_name__  ,			__magic_name__  ,			__magic_name__  ,			__magic_name__    )

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

							return optimizer, lr_scheduler







def     __lowerCamelCase							(   __magic_name__						:	Tuple  ,			__magic_name__						:	Dict    ):
							# 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

							a__:							Optional[int]									=sorted(glob.glob(F"{checkpoint_path}/global_step*"    )    )

							if len(__magic_name__    ) > 0:
														logger.info(F"Attempting to resume from {checkpoint_path}"    )
														# this magically updates self.optimizer and self.lr_scheduler
														a__				,      a__:							Union[str, Any]									=deepspeed_engine.load_checkpoint(
														    __magic_name__  ,			load_optimizer_states=__magic_name__  ,			load_lr_scheduler_states=__magic_name__    )
														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 shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
		from transformers import WhisperFeatureExtractor, WhisperProcessor


__UpperCAmelCase    	=  5_03_58
__UpperCAmelCase    	=  5_03_62


@require_torch
@require_torchaudio
@require_sentencepiece
class 						lowerCamelCase__			(		unittest.TestCase ):


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							str									="openai/whisper-small.en"
													a__:							Union[str, Any]									=tempfile.mkdtemp()


						def 	_lowerCamelCase      (			self    :	Optional[Any]   ,							**_a    :	Optional[int]      ):
													return WhisperTokenizer.from_pretrained(self.checkpoint   ,							**_a      )


						def 	_lowerCamelCase      (			self    :	str   ,							**_a    :	Tuple      ):
													return WhisperFeatureExtractor.from_pretrained(self.checkpoint   ,							**_a      )


						def 	_lowerCamelCase      (			self    :	Union[str, Any]      ):
													shutil.rmtree(self.tmpdirname      )


						def 	_lowerCamelCase      (			self    :	int      ):
													a__:							str									=self.get_tokenizer()
													a__:							Union[str, Any]									=self.get_feature_extractor()

													a__:							int									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													processor.save_pretrained(self.tmpdirname      )
													a__:							List[Any]									=WhisperProcessor.from_pretrained(self.tmpdirname      )

													self.assertEqual(processor.tokenizer.get_vocab()   ,							tokenizer.get_vocab()      )
													self.assertIsInstance(processor.tokenizer   ,							_a      )

													self.assertEqual(processor.feature_extractor.to_json_string()   ,							feature_extractor.to_json_string()      )
													self.assertIsInstance(processor.feature_extractor   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Dict      ):
													a__:							str									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													processor.save_pretrained(self.tmpdirname      )

													a__:							Tuple									=self.get_tokenizer(bos_token="(BOS)"   ,							eos_token="(EOS)"      )
													a__:							Optional[int]									=self.get_feature_extractor(do_normalize=_a   ,							padding_value=1.0      )

													a__:							Any									=WhisperProcessor.from_pretrained(
													    self.tmpdirname   ,							bos_token="(BOS)"   ,							eos_token="(EOS)"   ,							do_normalize=_a   ,							padding_value=1.0      )

													self.assertEqual(processor.tokenizer.get_vocab()   ,							tokenizer_add_kwargs.get_vocab()      )
													self.assertIsInstance(processor.tokenizer   ,							_a      )

													self.assertEqual(processor.feature_extractor.to_json_string()   ,							feature_extractor_add_kwargs.to_json_string()      )
													self.assertIsInstance(processor.feature_extractor   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Tuple      ):
													a__:							Optional[int]									=self.get_feature_extractor()
													a__:							str									=self.get_tokenizer()

													a__:							int									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[str]									=floats_list((3, 1_0_0_0)      )

													a__:							Optional[int]									=feature_extractor(_a   ,							return_tensors="np"      )
													a__:							Optional[int]									=processor(_a   ,							return_tensors="np"      )

													for key in input_feat_extract.keys():
																				self.assertAlmostEqual(input_feat_extract[key].sum()   ,							input_processor[key].sum()   ,							delta=1e-2      )


						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													a__:							Dict									=self.get_feature_extractor()
													a__:							int									=self.get_tokenizer()

													a__:							List[Any]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[Any]									="This is a test string"

													a__:							Tuple									=processor(text=_a      )

													a__:							int									=tokenizer(_a      )

													for key in encoded_tok.keys():
																				self.assertListEqual(encoded_tok[key]   ,							encoded_processor[key]      )


						def 	_lowerCamelCase      (			self    :	Any      ):
													a__:							int									=self.get_feature_extractor()
													a__:							List[str]									=self.get_tokenizer()

													a__:							List[str]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[Any]									=[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

													a__:							List[Any]									=processor.batch_decode(_a      )
													a__:							int									=tokenizer.batch_decode(_a      )

													self.assertListEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							List[Any]									=self.get_feature_extractor()
													a__:							Any									=self.get_tokenizer()

													a__:							Optional[int]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													self.assertListEqual(
													    processor.model_input_names   ,							feature_extractor.model_input_names   ,							msg="`processor` and `feature_extractor` model input names do not match"   ,							)


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							Any									=self.get_feature_extractor()
													a__:							str									=self.get_tokenizer()

													a__:							Dict									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )
													a__:							List[Any]									=processor.get_decoder_prompt_ids(task="transcribe"   ,							no_timestamps=_a      )

													self.assertIsInstance(_a   ,							_a      )
													for ids in forced_decoder_ids:
																				self.assertIsInstance(_a   ,							(list, tuple)      )

													a__:							Any									=[TRANSCRIBE, NOTIMESTAMPS]
													self.assertListEqual([ids[-1] for ids in forced_decoder_ids]   ,							_a      )


						def 	_lowerCamelCase      (			self    :	str      ):
													a__:							Tuple									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													a__:							List[Any]									=processor.get_prompt_ids("Mr. Quilter"      )
													a__:							Optional[Any]									=processor.tokenizer.decode(_a      )

													self.assertListEqual(prompt_ids.tolist()   ,							[5_0_3_6_0, 1_7_7_0, 1_3, 2_2_6_4, 3_4_6, 3_5_3]      )
													self.assertEqual(_a   ,							"<|startofprev|> Mr. Quilter"      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							int									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													a__:							List[str]									=processor.get_prompt_ids(""      )
													a__:							Any									=processor.tokenizer.decode(_a      )

													self.assertListEqual(prompt_ids.tolist()   ,							[5_0_3_6_0, 2_2_0]      )
													self.assertEqual(_a   ,							"<|startofprev|> "      )







						def 	_lowerCamelCase      (			self    :	Any      ):
													a__:							Tuple									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )

													def _test_prompt_error_raised_helper(_a    :	Union[str, Any]   ,							_a    :	List[str]      ):
																				with pytest.raises(_a      ) as excinfo:
																											processor.get_prompt_ids(_a      )
																				a__:							List[Any]									=F"Encountered text in the prompt corresponding to disallowed special token: {special_token}."
																				self.assertEqual(_a   ,							str(excinfo.value      )      )

													_test_prompt_error_raised_helper("<|startofprev|> test"   ,							"<|startofprev|>"      )
													_test_prompt_error_raised_helper("test <|notimestamps|>"   ,							"<|notimestamps|>"      )
													_test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"   ,							"<|zh|>"      )





