'''simple docstring'''




from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
						import tensorflow as tf

						from transformers import TFBlipTextModel
						from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST





class     __magic_name__						:







			def __init__(	self				:  int				,			lowercase_				:  int				,			lowercase_				:  Union[str, Any]=12				,			lowercase_				:  int=7				,			lowercase_				:  Optional[Any]=True				,			lowercase_				:  str=True				,			lowercase_				:  Union[str, Any]=True				,			lowercase_				:  Tuple=99				,			lowercase_				:  Dict=32				,			lowercase_				:  int=32				,			lowercase_				:  Optional[int]=2				,			lowercase_				:  Any=4				,			lowercase_				:  List[Any]=37				,			lowercase_				:  List[Any]=0.1				,			lowercase_				:  Any=0.1				,			lowercase_				:  str=512				,			lowercase_				:  Union[str, Any]=0.02				,			lowercase_				:  Optional[int]=0				,			lowercase_				:  List[Any]=None				,			):
										lowercase_					:			List[Any]										=       parent
										lowercase_					:			Union[str, Any]										=       batch_size
										lowercase_					:			Tuple										=       seq_length
										lowercase_					:			int										=       is_training
										lowercase_					:			Optional[int]										=       use_input_mask
										lowercase_					:			List[Any]										=       use_labels
										lowercase_					:			Tuple										=       vocab_size
										lowercase_					:			Optional[Any]										=       hidden_size
										lowercase_					:			int										=       projection_dim
										lowercase_					:			Tuple										=       num_hidden_layers
										lowercase_					:			int										=       num_attention_heads
										lowercase_					:			Union[str, Any]										=       intermediate_size
										lowercase_					:			Optional[Any]										=       dropout
										lowercase_					:			List[Any]										=       attention_dropout
										lowercase_					:			Optional[int]										=       max_position_embeddings
										lowercase_					:			Optional[Any]										=       initializer_range
										lowercase_					:			Optional[int]										=       scope
										lowercase_					:			List[Any]										=       bos_token_id







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str    ):
										lowercase_					:			Tuple										=       ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size    )

										lowercase_					:			List[Any]										=       None
										if self.use_input_mask:
																	lowercase_					:			Optional[int]										=       random_attention_mask([self.batch_size, self.seq_length]    )

										if input_mask is not None:
																	lowercase_					:			Tuple										=       input_mask.numpy()
																	lowercase_						,						lowercase_					:			Tuple										=       input_mask.shape
																	lowercase_					:			List[Any]										=       np.random.randint(1				,			seq_length - 1				,			size=(batch_size,)    )
																	for batch_idx, start_index in enumerate(lowercase_    ):
																								lowercase_					:			Dict										=       1
																								lowercase_					:			int										=       0

										lowercase_					:			Any										=       self.get_config()

										return config, input_ids, tf.convert_to_tensor(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										return BlipTextConfig(
										    vocab_size=self.vocab_size				,			hidden_size=self.hidden_size				,			projection_dim=self.projection_dim				,			num_hidden_layers=self.num_hidden_layers				,			num_attention_heads=self.num_attention_heads				,			intermediate_size=self.intermediate_size				,			dropout=self.dropout				,			attention_dropout=self.attention_dropout				,			max_position_embeddings=self.max_position_embeddings				,			initializer_range=self.initializer_range				,			bos_token_id=self.bos_token_id				,			)







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int				,			lowercase_				:  int				,			lowercase_				:  Dict				,			lowercase_				:  int    ):
										lowercase_					:			Optional[int]										=       TFBlipTextModel(config=lowercase_    )
										lowercase_					:			Tuple										=       model(lowercase_				,			attention_mask=lowercase_				,			training=lowercase_    )
										lowercase_					:			Optional[Any]										=       model(lowercase_				,			training=lowercase_    )
										self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)    )
										self.parent.assertEqual(result.pooler_output.shape				,			(self.batch_size, self.hidden_size)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Dict										=       self.prepare_config_and_inputs()
										lowercase_						,						lowercase_						,						lowercase_					:			Union[str, Any]										=       config_and_inputs
										lowercase_					:			Any										=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
										return config, inputs_dict





@require_tf
class     __magic_name__						(    _UpperCAmelCase,     unittest.TestCase):
			UpperCamelCase__				=       (TFBlipTextModel,) if is_tf_available() else ()
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Dict										=       BlipTextModelTester(self    )
										lowercase_					:			List[str]										=       ConfigTester(self				,			config_class=lowercase_				,			hidden_size=37    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										self.config_tester.run_common_tests()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			Optional[Any]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										pass







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										pass







			@unittest.skip(reason="""Blip does not use inputs_embeds"""    )
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										pass







			@unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING"""    )
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										pass







			@unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING"""    )
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										pass







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_					:			int										=       TFBlipTextModel.from_pretrained(lowercase_    )
																	self.assertIsNotNone(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  Tuple=True    ):
										super().test_pt_tf_model_equivalence(allow_missing_keys=lowercase_    )






'''simple docstring'''




import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler





class     __magic_name__						:







			def __init__(	self				:  str				,			lowercase_				:  Dict    ):
										if isinstance(lowercase_				,			lowercase_    ):
																	# Don't modify user's data should they want to reuse it (e.g. in tests), because once we
																	# modified it, it will not be accepted here again, since `auto` values would have been overridden
																	lowercase_					:			List[Any]										=       deepcopy(lowercase_    )
										elif os.path.exists(lowercase_    ):
																	with io.open(lowercase_				,			"""r"""				,			encoding="""utf-8"""    ) as f:
																								lowercase_					:			Union[str, Any]										=       json.load(lowercase_    )
										else:
																	try:
																								lowercase_					:			int										=       baseaa.urlsafe_baadecode(lowercase_    ).decode("""utf-8"""    )
																								lowercase_					:			str										=       json.loads(lowercase_    )
																	except (UnicodeDecodeError, AttributeError, ValueError):
																								raise ValueError(
																								    f'''Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}'''    )

										lowercase_					:			Any										=       config

										self.set_stage_and_offload()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										# zero stage - this is done as early as possible, before model is created, to allow
										# ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
										# during ``zero.Init()`` which needs to know the dtype, and some other hparams.
										lowercase_					:			Tuple										=       self.get_value("""zero_optimization.stage"""				,			-1    )

										# offload
										lowercase_					:			str										=       False
										if self.is_zeroa() or self.is_zeroa():
																	lowercase_					:			Dict										=       set(["""cpu""", """nvme"""]    )
																	lowercase_					:			List[Any]										=       set(
																	    [
																	        self.get_value("""zero_optimization.offload_optimizer.device"""    ),
																	        self.get_value("""zero_optimization.offload_param.device"""    ),
																	    ]    )
																	if len(offload_devices & offload_devices_valid    ) > 0:
																								lowercase_					:			Tuple										=       True







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  Any    ):
										lowercase_					:			Optional[Any]										=       self.config

										# find the config node of interest if it exists
										lowercase_					:			Tuple										=       ds_key_long.split("""."""    )
										lowercase_					:			Union[str, Any]										=       nodes.pop()
										for node in nodes:
																	lowercase_					:			List[str]										=       config.get(lowercase_    )
																	if config is None:
																								return None, ds_key

										return config, ds_key







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str				,			lowercase_				:  List[str]				,			lowercase_				:  List[str]=None    ):
										lowercase_						,						lowercase_					:			List[Any]										=       self.find_config_node(lowercase_    )
										if config is None:
																	return default
										return config.get(lowercase_				,			lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]				,			lowercase_				:  Optional[int]				,			lowercase_				:  int=False    ):
										lowercase_					:			int										=       self.config

										# find the config node of interest if it exists
										lowercase_					:			Dict										=       ds_key_long.split("""."""    )
										for node in nodes:
																	lowercase_					:			List[Any]										=       config
																	lowercase_					:			Dict										=       config.get(lowercase_    )
																	if config is None:
																								if must_exist:
																															raise ValueError(f'''Can\'t find {ds_key_long} entry in the config: {self.config}'''    )
																								else:
																															return

        # if found remove it
										if parent_config is not None:
																	parent_config.pop(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]				,			lowercase_				:  Tuple    ):
										lowercase_					:			str										=       self.get_value(lowercase_    )
										return False if value is None else bool(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]				,			lowercase_				:  Union[str, Any]    ):
										lowercase_					:			Union[str, Any]										=       self.get_value(lowercase_    )
										return False if value is None else not bool(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										return self._stage == 2







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										return self._stage == 3







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										return self._offload





class     __magic_name__						:







			def __init__(	self				:  Any				,			lowercase_				:  Union[str, Any]    ):
										lowercase_					:			Any										=       engine







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  int				,			**lowercase_				:  str    ):
										# runs backpropagation and handles mixed precision
										self.engine.backward(lowercase_				,			**lowercase_    )

										# Deepspeed's `engine.step` performs the following operations:
										# - gradient accumulation check
										# - gradient clipping
										# - optimizer step
										# - zero grad
										# - checking overflow
										# - lr_scheduler step (only if engine.lr_scheduler is not None)
										self.engine.step()
										# and this plugin overrides the above calls with no-ops when Accelerate runs under
										# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
										# training loop that works transparently under many training regimes.





class     __magic_name__						(    _UpperCAmelCase):







			def __init__(	self				:  Optional[Any]				,			lowercase_				:  Tuple    ):
										super().__init__(lowercase_				,			device_placement=lowercase_				,			scaler=lowercase_    )
										lowercase_					:			Any										=       hasattr(self.optimizer				,			"""overflow"""    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int				,			lowercase_				:  Tuple=None    ):
										pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										if self.__has_overflow__:
																	return self.optimizer.overflow
										return False





class     __magic_name__						(    _UpperCAmelCase):







			def __init__(	self				:  Dict				,			lowercase_				:  Optional[Any]				,			lowercase_				:  Tuple    ):
										super().__init__(lowercase_				,			lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed





class     __magic_name__						:







			def __init__(	self				:  Any				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  List[str]=0.0_01				,			lowercase_				:  List[str]=0				,			**lowercase_				:  List[Any]    ):
										lowercase_					:			str										=       params
										lowercase_					:			List[Any]										=       lr
										lowercase_					:			int										=       weight_decay
										lowercase_					:			Union[str, Any]										=       kwargs





class     __magic_name__						:







			def __init__(	self				:  Tuple				,			lowercase_				:  Optional[Any]				,			lowercase_				:  List[str]=None				,			lowercase_				:  int=0				,			**lowercase_				:  int    ):
										lowercase_					:			Union[str, Any]										=       optimizer
										lowercase_					:			List[str]										=       total_num_steps
										lowercase_					:			Dict										=       warmup_num_steps
										lowercase_					:			Dict										=       kwargs





