





'''simple docstring'''







import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask


if is_torch_available():
					import torch

					from transformers import BlipTextModel
					from transformers.models.blip.modeling_blip import BLIP_PRETRAINED_MODEL_ARCHIVE_LIST


class        __a   :

							def __init__(				self				:  Optional[Any]							,  __magic_name__				:  List[Any]							,  __magic_name__				:  List[Any]=12							,  __magic_name__				:  Optional[int]=7							,  __magic_name__				:  int=True							,  __magic_name__				:  Union[str, Any]=True							,  __magic_name__				:  Union[str, Any]=True							,  __magic_name__				:  int=99							,  __magic_name__				:  List[str]=32							,  __magic_name__				:  Tuple=32							,  __magic_name__				:  str=5							,  __magic_name__				:  Union[str, Any]=4							,  __magic_name__				:  Optional[int]=37							,  __magic_name__				:  List[Any]=0.1							,  __magic_name__				:  Optional[Any]=0.1							,  __magic_name__				:  Optional[Any]=5_12							,  __magic_name__				:  Any=0.0_2							,  __magic_name__				:  int=0							,  __magic_name__				:  Any=None							,  )       ->							List[str]:







													"""simple docstring"""
													UpperCAmelCase_      :   Union[str, Any]      =	parent
													UpperCAmelCase_      :   List[Any]      =	batch_size
													UpperCAmelCase_      :   Optional[int]      =	seq_length
													UpperCAmelCase_      :   Any      =	is_training
													UpperCAmelCase_      :   Optional[Any]      =	use_input_mask
													UpperCAmelCase_      :   List[str]      =	use_labels
													UpperCAmelCase_      :   List[Any]      =	vocab_size
													UpperCAmelCase_      :   str      =	hidden_size
													UpperCAmelCase_      :   Dict      =	projection_dim
													UpperCAmelCase_      :   Dict      =	num_hidden_layers
													UpperCAmelCase_      :   Any      =	num_attention_heads
													UpperCAmelCase_      :   Union[str, Any]      =	intermediate_size
													UpperCAmelCase_      :   Optional[Any]      =	dropout
													UpperCAmelCase_      :   Any      =	attention_dropout
													UpperCAmelCase_      :   Any      =	max_position_embeddings
													UpperCAmelCase_      :   Tuple      =	initializer_range
													UpperCAmelCase_      :   Optional[int]      =	scope
													UpperCAmelCase_      :   Dict      =	bos_token_id

							def      UpperCAmelCase__						(				self				:  Optional[int] )       ->							List[Any]:







													"""simple docstring"""
													UpperCAmelCase_      :   List[Any]      =	ids_tensor([self.batch_size, self.seq_length]							,  self.vocab_size )

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

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

													UpperCAmelCase_      :   int      =	self.get_config()

													return config, input_ids, input_mask

							def      UpperCAmelCase__						(				self				:  Any )       ->							Tuple:







													"""simple docstring"""
													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      UpperCAmelCase__						(				self				:  Any							,  __magic_name__				:  Any							,  __magic_name__				:  List[str]							,  __magic_name__				:  Optional[int] )       ->							Dict:







													"""simple docstring"""
													UpperCAmelCase_      :   List[Any]      =	BlipTextModel(config=__magic_name__ )
													model.to(__magic_name__ )
													model.eval()
													with torch.no_grad():
																			UpperCAmelCase_      :   Tuple      =	model(__magic_name__							,  attention_mask=__magic_name__ )
																			UpperCAmelCase_      :   Any      =	model(__magic_name__ )
													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      UpperCAmelCase__						(				self				:  Union[str, Any] )       ->							List[str]:







													"""simple docstring"""
													UpperCAmelCase_      :   Dict      =	self.prepare_config_and_inputs()
													UpperCAmelCase_      :   Any      =	config_and_inputs
													UpperCAmelCase_      :   Optional[int]      =	{"""input_ids""": input_ids, """attention_mask""": input_mask}
													return config, inputs_dict





@require_torch
class        __a   (lowerCamelCase							,		unittest.TestCase  ):
							__a       : Tuple          =							(BlipTextModel,) if is_torch_available() else ()
							__a       : Any          =							False
							__a       : Tuple          =							False
							__a       : Optional[Any]          =							False

							def      UpperCAmelCase__						(				self				:  Any )       ->							Dict:







													"""simple docstring"""
													UpperCAmelCase_      :   int      =	BlipTextModelTester(self )
													UpperCAmelCase_      :   int      =	ConfigTester(self							,  config_class=__magic_name__							,  hidden_size=37 )

							def      UpperCAmelCase__						(				self				:  str )       ->							List[str]:







													"""simple docstring"""
													self.config_tester.run_common_tests()

							def      UpperCAmelCase__						(				self				:  Tuple )       ->							List[Any]:







													"""simple docstring"""
													UpperCAmelCase_      :   Any      =	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*__magic_name__ )

							def      UpperCAmelCase__						(				self				:  List[Any] )       ->							Optional[int]:







													"""simple docstring"""
													pass

							def      UpperCAmelCase__						(				self				:  Dict )       ->							List[str]:







													"""simple docstring"""
													pass

							@unittest.skip(reason='''Blip does not use inputs_embeds''' )
							def      UpperCAmelCase__						(				self				:  Tuple )       ->							str:







													"""simple docstring"""
													pass

							@unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING''' )
							def      UpperCAmelCase__						(				self				:  Optional[int] )       ->							List[str]:







													"""simple docstring"""
													pass

							@unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING''' )
							def      UpperCAmelCase__						(				self				:  List[str] )       ->							Optional[int]:







													"""simple docstring"""
													pass

							@slow
							def      UpperCAmelCase__						(				self				:  Optional[Any] )       ->							str:







													"""simple docstring"""





													for model_name in BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			UpperCAmelCase_      :   Dict      =	BlipTextModel.from_pretrained(__magic_name__ )
																			self.assertIsNotNone(__magic_name__ )



							def      UpperCAmelCase__						(				self				:  Union[str, Any] )       ->							List[str]:







													"""simple docstring"""
													super().test_pt_tf_model_equivalence(allow_missing_keys=__magic_name__ )













"""simple docstring"""



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


A_     =				get_logger(__name__)




def   _lowerCAmelCase		(				UpperCAmelCase__							:       str,    UpperCAmelCase__							:       List[dict],    UpperCAmelCase__							:       list,    UpperCAmelCase__							:       list,    UpperCAmelCase__							:       int,    UpperCAmelCase__							:       GradScaler = None,    )		->Any:
						# Model states
						for i, state in enumerate(UpperCAmelCase__   ):
												A__    :	Optional[int]			=	f'{MODEL_NAME}.bin' if i == 0 else f'{MODEL_NAME}_{i}.bin'
												A__    :	int			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Model weights saved in {output_model_file}'   )
						# Optimizer states
						for i, opt in enumerate(UpperCAmelCase__   ):
												A__    :	str			=	opt.state_dict()
												A__    :	Any			=	f'{OPTIMIZER_NAME}.bin' if i == 0 else f'{OPTIMIZER_NAME}_{i}.bin'
												A__    :	Optional[int]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Optimizer state saved in {output_optimizer_file}'   )
						# Scheduler states
						for i, scheduler in enumerate(UpperCAmelCase__   ):
												A__    :	Union[str, Any]			=	scheduler.state_dict()
												A__    :	Union[str, Any]			=	f'{SCHEDULER_NAME}.bin' if i == 0 else f'{SCHEDULER_NAME}_{i}.bin'
												A__    :	Tuple			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Scheduler state saved in {output_scheduler_file}'   )
						# GradScaler state
						if scaler is not None:
												A__    :	List[str]			=	scaler.state_dict()
												A__    :	Optional[Any]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												torch.save(UpperCAmelCase__,    UpperCAmelCase__   )
												logger.info(f'Gradient scaler state saved in {output_scaler_file}'   )
						# Random number generator states
						A__    :	str			=	{}
						A__    :	str			=	f'{RNG_STATE_NAME}_{process_index}.pkl'
						A__    :	Optional[Any]			=	random.getstate()
						A__    :	str			=	np.random.get_state()
						A__    :	List[str]			=	torch.get_rng_state()
						if is_xpu_available():
												A__    :	Optional[Any]			=	torch.xpu.get_rng_state_all()
						else:
												A__    :	str			=	torch.cuda.get_rng_state_all()
						if is_tpu_available():
												A__    :	Tuple			=	xm.get_rng_state()
						A__    :	str			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
						torch.save(UpperCAmelCase__,    UpperCAmelCase__   )
						logger.info(f'Random states saved in {output_states_file}'   )
						return output_dir




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Any,    UpperCAmelCase__							:       Union[str, Any],    UpperCAmelCase__							:       Union[str, Any],    UpperCAmelCase__							:       Dict,    UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       Optional[Any]=None,    UpperCAmelCase__							:       str=None,    **UpperCAmelCase__							:       Optional[int],    )		->List[str]:
						if map_location not in [None, "cpu", "on_device"]:
												raise TypeError(
												    """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""   )
						if map_location is None:
												A__    :	Union[str, Any]			=	"""cpu"""
						elif map_location == "on_device":
												A__    :	Any			=	PartialState().device
						# Model states
						for i, model in enumerate(UpperCAmelCase__   ):
												A__    :	Any			=	f'{MODEL_NAME}.bin' if i == 0 else f'{MODEL_NAME}_{i}.bin'
												A__    :	List[str]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												models[i].load_state_dict(torch.load(UpperCAmelCase__,    map_location=UpperCAmelCase__   ),    **UpperCAmelCase__   )
						logger.info("""All model weights loaded successfully"""   )

						# Optimizer states
						for i, opt in enumerate(UpperCAmelCase__   ):
												A__    :	Union[str, Any]			=	f'{OPTIMIZER_NAME}.bin' if i == 0 else f'{OPTIMIZER_NAME}_{i}.bin'
												A__    :	Dict			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												A__    :	List[Any]			=	torch.load(UpperCAmelCase__,    map_location=UpperCAmelCase__   )
												optimizers[i].load_state_dict(UpperCAmelCase__   )
						logger.info("""All optimizer states loaded successfully"""   )

						# Scheduler states
						for i, scheduler in enumerate(UpperCAmelCase__   ):
												A__    :	List[Any]			=	f'{SCHEDULER_NAME}.bin' if i == 0 else f'{SCHEDULER_NAME}_{i}.bin'
												A__    :	Optional[Any]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												scheduler.load_state_dict(torch.load(UpperCAmelCase__   )   )
						logger.info("""All scheduler states loaded successfully"""   )

						# GradScaler state
						if scaler is not None:
												A__    :	List[str]			=	os.path.join(UpperCAmelCase__,    UpperCAmelCase__   )
												scaler.load_state_dict(torch.load(UpperCAmelCase__   )   )
												logger.info("""GradScaler state loaded successfully"""   )

						# Random states
						try:
												A__    :	Any			=	torch.load(os.path.join(UpperCAmelCase__,    f'{RNG_STATE_NAME}_{process_index}.pkl'   )   )
												random.setstate(states["""random_state"""]   )
												np.random.set_state(states["""numpy_random_seed"""]   )
												torch.set_rng_state(states["""torch_manual_seed"""]   )
												if is_xpu_available():
																		torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]   )
												else:
																		torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]   )
												if is_tpu_available():
																		xm.set_rng_state(states["""xm_seed"""]   )
												logger.info("""All random states loaded successfully"""   )
						except Exception:
												logger.info("""Could not load random states"""   )




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       Dict,    UpperCAmelCase__							:       int = 0   )		->Dict:
						A__    :	Dict			=	Path(UpperCAmelCase__   ) / f'custom_checkpoint_{index}.pkl'
						logger.info(f'Saving the state of {get_pretty_name(UpperCAmelCase__   )} to {save_location}'   )
						torch.save(obj.state_dict(),    UpperCAmelCase__   )




def   _lowerCAmelCase		(				UpperCAmelCase__							:       Tuple,    UpperCAmelCase__							:       str,    UpperCAmelCase__							:       int = 0   )		->str:
						A__    :	Optional[int]			=	f'{path}/custom_checkpoint_{index}.pkl'
						logger.info(f'Loading the state of {get_pretty_name(UpperCAmelCase__   )} from {load_location}'   )
						obj.load_state_dict(torch.load(UpperCAmelCase__,    map_location="""cpu"""   )   )






