




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    __lowerCAmelCase       :






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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         lowercase    :List[str]      	=						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    :Optional[Any]      	=						input_mask.numpy()
           lowercase    :Optional[Any]      	=						input_mask.shape
           lowercase    :Optional[int]      	=						np.random.randint(1       ,					seq_length - 1       ,					size=(batch_size,)				)
           for batch_idx, start_index in enumerate(_lowerCAmelCase				):
             lowercase    :str      	=						1
             lowercase    :List[str]      	=						0

         lowercase    :List[Any]      	=						self.get_config()

         return config, input_ids, tf.convert_to_tensor(_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         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:					List[str]       ,					_lowerCAmelCase:					List[str]       ,					_lowerCAmelCase:					List[Any]       ,					_lowerCAmelCase:					Optional[int]				):
         lowercase    :Dict      	=						TFBlipTextModel(config=_lowerCAmelCase				)
         lowercase    :str      	=						model(_lowerCAmelCase       ,					attention_mask=_lowerCAmelCase       ,					training=_lowerCAmelCase				)
         lowercase    :str      	=						model(_lowerCAmelCase       ,					training=_lowerCAmelCase				)
         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    :Any      	=						self.prepare_config_and_inputs()
         lowercase    :int      	=						config_and_inputs
         lowercase    :str      	=						{"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict







@require_tf
class    __lowerCAmelCase       (							lowerCAmelCase      , unittest.TestCase):
       _a 					=						(TFBlipTextModel,) if is_tf_available() else ()
       _a 					=						False
       _a 					=						False
       _a 					=						False






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Tuple      	=						BlipTextModelTester(self				)
         lowercase    :str      	=						ConfigTester(self       ,					config_class=_lowerCAmelCase       ,					hidden_size=37				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         self.config_tester.run_common_tests()






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :List[Any]      	=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCAmelCase				)






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






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         pass






       @unittest.skip(reason="Blip does not use inputs_embeds"				)
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         pass






       @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"				)
       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:					Union[str, Any]				):
         pass






       @slow
       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase    :Dict      	=						TFBlipTextModel.from_pretrained(_lowerCAmelCase				)
           self.assertIsNotNone(_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict       ,					_lowerCAmelCase:					Optional[int]=True				):
         super().test_pt_tf_model_equivalence(allow_missing_keys=_lowerCAmelCase				)









import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment



def       UpperCAmelCase__  (     ):
  lowercase    :List[str]      	=						torch.nn.Linear(2,			4   )
  lowercase    :List[Any]      	=						torch.optim.AdamW(model.parameters(),			lr=1.0   )
  lowercase    :int      	=						torch.optim.lr_scheduler.OneCycleLR(lowerCamelCase,			max_lr=0.01,			steps_per_epoch=2,			epochs=1   )
  lowercase    :Optional[int]      	=						DataLoader(TensorDataset(torch.tensor([1, 2, 3]   )   )   )
  lowercase    :List[str]      	=						DataLoader(TensorDataset(torch.tensor([4, 5, 6]   )   )   )

  return model, optimizer, scheduler, train_dl, valid_dl



def       UpperCAmelCase__  (     lowerCamelCase   ):
  return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Dict      	=						torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
  model.load_state_dict(lowerCamelCase   )







class    __lowerCAmelCase       (							lowerCAmelCase):






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :List[str]      	=						Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Any      	=						Accelerator(cpu=_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :Any      	=						Accelerator()
         lowercase    :Dict      	=						GradientState()
         assert state.num_steps == 1
         lowercase    :Tuple      	=						4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         lowercase    :List[Any]      	=						False
         assert state.sync_gradients is False
         GradientState._reset_state()






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :str      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Any      	=						create_components()

         (
             (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    
         )    :Tuple      	=						accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         self.assertTrue(prepared_model in accelerator._models				)
         self.assertTrue(prepared_optimizer in accelerator._optimizers				)
         self.assertTrue(prepared_scheduler in accelerator._schedulers				)
         self.assertTrue(prepared_train_dl in accelerator._dataloaders				)
         self.assertTrue(prepared_valid_dl in accelerator._dataloaders				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         accelerator.free_memory()

         self.assertTrue(len(accelerator._models				) == 0				)
         self.assertTrue(len(accelerator._optimizers				) == 0				)
         self.assertTrue(len(accelerator._schedulers				) == 0				)
         self.assertTrue(len(accelerator._dataloaders				) == 0				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         PartialState._reset_state()

         # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
         def noop(*_lowerCAmelCase:					List[str]       ,					**_lowerCAmelCase:					Optional[int]				):
           pass

         with patch("torch.cuda.set_device"       ,					_lowerCAmelCase				), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"				):
           lowercase    :List[Any]      	=						Accelerator()
           self.assertEqual(str(accelerator.state.device				)       ,					"cuda:64"				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :Tuple      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Optional[Any]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :Tuple      	=						get_signature(_lowerCAmelCase				)

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # make sure loaded weights match
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :int      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :int      	=						get_signature(_lowerCAmelCase				)

         # saving hook
         def save_config(_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					str       ,					_lowerCAmelCase:					List[Any]				):
           lowercase    :Dict      	=						{"class_name": models[0].__class__.__name__}

           with open(os.path.join(_lowerCAmelCase       ,					"data.json"				)       ,					"w"				) as f:
             json.dump(_lowerCAmelCase       ,					_lowerCAmelCase				)

        # loading hook
         def load_config(_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					Union[str, Any]				):
           with open(os.path.join(_lowerCAmelCase       ,					"data.json"				)       ,					"r"				) as f:
             lowercase    :int      	=						json.load(_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						config["class_name"]

         lowercase    :Optional[Any]      	=						accelerator.register_save_state_pre_hook(_lowerCAmelCase				)
         lowercase    :Tuple      	=						accelerator.register_load_state_pre_hook(_lowerCAmelCase				)

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match with hooks
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # random class name to verify correct one is loaded
           lowercase    :Optional[int]      	=						"random"

           # make sure loaded weights match with hooks
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)

           # mode.class_name is loaded from config
           self.assertTrue(model.class_name == model.__class__.__name__				)

         # remove hooks
         save_hook.remove()
         load_hook.remove()

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match with hooks removed
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # random class name to verify correct one is loaded
           lowercase    :str      	=						"random"

           # make sure loaded weights match with hooks removed
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)

           # mode.class_name is NOT loaded from config
           self.assertTrue(model.class_name != model.__class__.__name__				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :List[str]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[Any]      	=						create_components()
         lowercase    :List[str]      	=						None

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :str      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertTrue(dummy_obj is None				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Union[str, Any]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Union[str, Any]      	=						create_components()
         lowercase    :Union[str, Any]      	=						[1, 2, 3]

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Dummy object should have `_is_accelerate_prepared` set to `True`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Model is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Optimizer is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Scheduler is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map={"": 0}       ,					)
         lowercase    :int      	=						Accelerator()

         # This should work
         lowercase    :Optional[int]      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Optional[Any]      	=						Accelerator()

         with init_empty_weights():
           lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Dict      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :int      	=						"cpu"

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					device_map=_lowerCAmelCase       ,					load_in_abit=_lowerCAmelCase       ,					llm_inta_enable_fpaa_cpu_offload=_lowerCAmelCase				)

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Dict      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         from transformers import AutoModelForCausalLM

         lowercase    :List[str]      	=						{"distributed_type": DistributedType.MULTI_GPU}

         with init_empty_weights():
           lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Optional[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Any      	=						1

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :List[str]      	=						Accelerator()

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						accelerator.prepare(_lowerCAmelCase				)

         PartialState._reset_state()






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         from transformers import AutoModelForCausalLM

         with init_empty_weights():
           lowercase    :List[str]      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           lowercase    :List[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Optional[int]      	=						1

         lowercase    :Union[str, Any]      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :Union[str, Any]      	=						Accelerator()

         # This should work
         lowercase    :List[Any]      	=						accelerator.prepare(_lowerCAmelCase				)






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Optional[int]      	=						torch.nn.Linear(10       ,					10				)
         lowercase    :Optional[int]      	=						torch.optim.SGD(model.parameters()       ,					lr=0.01				)
         lowercase    :List[Any]      	=						Accelerator(cpu=_lowerCAmelCase				)
         lowercase    :List[str]      	=						accelerator.prepare(_lowerCAmelCase				)



