
"""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      _UpperCAmelCase							:




      def __init__(							self   ,       lowercase_   ,       lowercase_=1_2   ,       lowercase_=7   ,       lowercase_=True   ,       lowercase_=True   ,       lowercase_=True   ,       lowercase_=9_9   ,       lowercase_=3_2   ,       lowercase_=3_2   ,       lowercase_=2   ,       lowercase_=4   ,       lowercase_=3_7   ,       lowercase_=0.1   ,       lowercase_=0.1   ,       lowercase_=5_1_2   ,       lowercase_=0.0_2   ,       lowercase_=0   ,       lowercase_=None   ,       )						->					List[str]:
          UpperCAmelCase									=					parent
          UpperCAmelCase									=					batch_size
          UpperCAmelCase									=					seq_length
          UpperCAmelCase									=					is_training
          UpperCAmelCase									=					use_input_mask
          UpperCAmelCase									=					use_labels
          UpperCAmelCase									=					vocab_size
          UpperCAmelCase									=					hidden_size
          UpperCAmelCase									=					projection_dim
          UpperCAmelCase									=					num_hidden_layers
          UpperCAmelCase									=					num_attention_heads
          UpperCAmelCase									=					intermediate_size
          UpperCAmelCase									=					dropout
          UpperCAmelCase									=					attention_dropout
          UpperCAmelCase									=					max_position_embeddings
          UpperCAmelCase									=					initializer_range
          UpperCAmelCase									=					scope
          UpperCAmelCase									=					bos_token_id




      def     a_      (							self							)						->					str:
          UpperCAmelCase									=					ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size							)

          UpperCAmelCase									=					None
          if self.use_input_mask:
              UpperCAmelCase									=					random_attention_mask([self.batch_size, self.seq_length]							)

          if input_mask is not None:
              UpperCAmelCase									=					input_mask.numpy()
              UpperCAmelCase			,							UpperCAmelCase									=					input_mask.shape
              UpperCAmelCase									=					np.random.randint(1   ,       seq_length - 1   ,       size=(batch_size,)							)
              for batch_idx, start_index in enumerate(A_							):
                  UpperCAmelCase									=					1
                  UpperCAmelCase									=					0

          UpperCAmelCase									=					self.get_config()

          return config, input_ids, tf.convert_to_tensor(A_							)




      def     a_      (							self							)						->					List[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     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					Dict:
          UpperCAmelCase									=					TFBlipTextModel(config=A_							)
          UpperCAmelCase									=					model(A_   ,       attention_mask=A_   ,       training=A_							)
          UpperCAmelCase									=					model(A_   ,       training=A_							)
          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     a_      (							self							)						->					Optional[Any]:
          UpperCAmelCase									=					self.prepare_config_and_inputs()
          UpperCAmelCase			,							UpperCAmelCase			,							UpperCAmelCase									=					config_and_inputs
          UpperCAmelCase									=					{'input_ids': input_ids, 'attention_mask': input_mask}
          return config, inputs_dict

@require_tf
class      _UpperCAmelCase							(		_lowercase     ,			unittest.TestCase							):
      __SCREAMING_SNAKE_CASE   :							Tuple							      =	(TFBlipTextModel,) if is_tf_available() else ()
      __SCREAMING_SNAKE_CASE   :							Any							      =	False
      __SCREAMING_SNAKE_CASE   :							Any							      =	False
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	False




      def     a_      (							self							)						->					Tuple:
          UpperCAmelCase									=					BlipTextModelTester(self							)
          UpperCAmelCase									=					ConfigTester(self   ,       config_class=A_   ,       hidden_size=3_7							)




      def     a_      (							self							)						->					str:
          self.config_tester.run_common_tests()




      def     a_      (							self							)						->					Dict:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*A_							)




      def     a_      (							self							)						->					int:
          pass




      def     a_      (							self							)						->					List[Any]:
          pass




      @unittest.skip(reason='Blip does not use inputs_embeds'							)
      def     a_      (							self							)						->					Any:
          pass




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




      @unittest.skip(reason='BlipTextModel has no base class and is not available in MODEL_MAPPING'							)
      def     a_      (							self							)						->					int:
          pass




      @slow
      def     a_      (							self							)						->					str:
          for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase									=					TFBlipTextModel.from_pretrained(A_							)
              self.assertIsNotNone(A_							)




      def     a_      (							self   ,       lowercase_=True							)						->					List[str]:
          super().test_pt_tf_model_equivalence(allow_missing_keys=A_							)






"""simple docstring"""

import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu
class      _UpperCAmelCase							(		nn.Module							):




      def __init__(							self							)						->					Any:
          super().__init__()
          UpperCAmelCase									=					nn.Linear(3   ,       4							)
          UpperCAmelCase									=					nn.BatchNormad(4							)
          UpperCAmelCase									=					nn.Linear(4   ,       5							)




      def     a_      (							self   ,       lowercase_							)						->					Any:
          return self.lineara(self.batchnorm(self.lineara(lowercase_							)							)							)

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):




      def     a_      (							self   ,       lowercase_   ,       *lowercase_   ,       **lowercase_							)						->					List[str]:
          return (args[0] + 1,) + args[1:], kwargs

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):




      def     a_      (							self   ,       lowercase_   ,       lowercase_							)						->					Union[str, Any]:
          return output + 1

class      _UpperCAmelCase							(		unittest.TestCase							):




      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					ModelHook()

          add_hook_to_module(lowercase_   ,       lowercase_							)
          self.assertEqual(test_model._hf_hook   ,       lowercase_							)
          self.assertTrue(hasattr(lowercase_   ,       '_old_forward'							)							)

          # Check adding the hook did not change the name or the signature
          self.assertEqual(test_model.forward.__name__   ,       'forward'							)
          self.assertListEqual(list(inspect.signature(test_model.forward							).parameters							)   ,       ['x']							)

          remove_hook_from_module(lowercase_							)
          self.assertFalse(hasattr(lowercase_   ,       '_hf_hook'							)							)
          self.assertFalse(hasattr(lowercase_   ,       '_old_forward'							)							)




      def     a_      (							self							)						->					Any:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					ModelHook()

          add_hook_to_module(lowercase_   ,       lowercase_							)
          add_hook_to_module(lowercase_   ,       lowercase_   ,       append=lowercase_							)

          self.assertEqual(isinstance(test_model._hf_hook   ,       lowercase_							)   ,       lowercase_							)
          self.assertEqual(len(test_model._hf_hook.hooks							)   ,       2							)
          self.assertTrue(hasattr(lowercase_   ,       '_old_forward'							)							)

          # Check adding the hook did not change the name or the signature
          self.assertEqual(test_model.forward.__name__   ,       'forward'							)
          self.assertListEqual(list(inspect.signature(test_model.forward							).parameters							)   ,       ['x']							)

          remove_hook_from_module(lowercase_							)
          self.assertFalse(hasattr(lowercase_   ,       '_hf_hook'							)							)
          self.assertFalse(hasattr(lowercase_   ,       '_old_forward'							)							)




      def     a_      (							self							)						->					Any:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(x + 1							)
          UpperCAmelCase									=					test_model(x + 2							)

          UpperCAmelCase									=					PreForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       lowercase_   ,       atol=1E-5							)							)

          # Attaching a hook to a model when it already has one replaces, does not chain
          UpperCAmelCase									=					PreForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       lowercase_   ,       atol=1E-5							)							)

          # You need to use the sequential hook to chain two or more hooks
          UpperCAmelCase									=					SequentialHook(PreForwardHook()   ,       PreForwardHook()							)
          add_hook_to_module(lowercase_   ,       lowercase_							)

          UpperCAmelCase									=					test_model(lowercase_							)
          assert torch.allclose(lowercase_   ,       lowercase_   ,       atol=1E-5							)




      def     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(lowercase_							)

          UpperCAmelCase									=					PostForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       output + 1   ,       atol=1E-5							)							)

          # Attaching a hook to a model when it already has one replaces, does not chain
          UpperCAmelCase									=					PostForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       output + 1   ,       atol=1E-5							)							)

          # You need to use the sequential hook to chain two or more hooks
          UpperCAmelCase									=					SequentialHook(PostForwardHook()   ,       PostForwardHook()							)
          add_hook_to_module(lowercase_   ,       lowercase_							)

          UpperCAmelCase									=					test_model(lowercase_							)
          assert torch.allclose(lowercase_   ,       output + 2   ,       atol=1E-5							)




      def     a_      (							self							)						->					Tuple:
          UpperCAmelCase									=					ModelForTest()
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					test_model(lowercase_							)

          UpperCAmelCase									=					PostForwardHook()
          add_hook_to_module(lowercase_   ,       lowercase_							)
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertTrue(torch.allclose(lowercase_   ,       output + 1							)							)
          self.assertTrue(outputa.requires_grad							)

          UpperCAmelCase									=					True
          UpperCAmelCase									=					test_model(lowercase_							)
          self.assertFalse(outputa.requires_grad							)




      @require_multi_gpu
      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					ModelForTest()
          # Everything is on CPU
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # This will move each submodule on different devices
          add_hook_to_module(model.lineara   ,       AlignDevicesHook(execution_device=0							)							)
          add_hook_to_module(model.batchnorm   ,       AlignDevicesHook(execution_device=0							)							)
          add_hook_to_module(model.lineara   ,       AlignDevicesHook(execution_device=1							)							)

          self.assertEqual(model.lineara.weight.device   ,       torch.device(0							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device(0							)							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       torch.device(0							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device(1							)							)

          # We can still make a forward pass. The input does not need to be on any particular device
          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       torch.device(1							)							)

          # We can add a general hook to put back output on same device as input.
          add_hook_to_module(lowercase_   ,       AlignDevicesHook(io_same_device=lowercase_							)							)
          UpperCAmelCase									=					torch.randn(2   ,       3							).to(0							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       torch.device(0							)							)




      def     a_      (							self							)						->					Dict:
          UpperCAmelCase									=					ModelForTest()

          # Everything is on CPU
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # This will move each submodule on different devices
          UpperCAmelCase									=					{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

          add_hook_to_module(model.lineara   ,       AlignDevicesHook(**lowercase_							)							)
          add_hook_to_module(model.batchnorm   ,       AlignDevicesHook(**lowercase_							)							)
          add_hook_to_module(model.lineara   ,       AlignDevicesHook(**lowercase_							)							)

          # Parameters have been offloaded, so on the meta device
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          # Buffers are not included in the offload by default, so are on the execution device
          UpperCAmelCase									=					torch.device(hook_kwargs['execution_device']							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_module(model.lineara							)
          remove_hook_from_module(model.batchnorm							)
          remove_hook_from_module(model.lineara							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # Now test with buffers included in the offload
          UpperCAmelCase									=					{
              'execution_device': 0 if torch.cuda.is_available() else 'cpu',
              'offload': True,
              'offload_buffers': True,
          }

          add_hook_to_module(model.lineara   ,       AlignDevicesHook(**lowercase_							)							)
          add_hook_to_module(model.batchnorm   ,       AlignDevicesHook(**lowercase_							)							)
          add_hook_to_module(model.lineara   ,       AlignDevicesHook(**lowercase_							)							)

          # Parameters have been offloaded, so on the meta device, buffers included
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       torch.device('meta'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_module(model.lineara							)
          remove_hook_from_module(model.batchnorm							)
          remove_hook_from_module(model.lineara							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)




      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					ModelForTest()

          # Everything is on CPU
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # This will move each submodule on different devices
          UpperCAmelCase									=					0 if torch.cuda.is_available() else 'cpu'
          attach_align_device_hook(lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_							)

          # Parameters have been offloaded, so on the meta device
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          # Buffers are not included in the offload by default, so are on the execution device
          UpperCAmelCase									=					torch.device(lowercase_							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # Now test with buffers included in the offload
          attach_align_device_hook(lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       offload_buffers=lowercase_							)

          # Parameters have been offloaded, so on the meta device, buffers included
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       torch.device('meta'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)




      def     a_      (							self							)						->					Dict:
          UpperCAmelCase									=					ModelForTest()

          # Everything is on CPU
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # This will move each submodule on different devices
          UpperCAmelCase									=					0 if torch.cuda.is_available() else 'cpu'
          attach_align_device_hook(
              lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       weights_map=model.state_dict()							)

          # Parameters have been offloaded, so on the meta device
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          # Buffers are not included in the offload by default, so are on the execution device
          UpperCAmelCase									=					torch.device(lowercase_							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       lowercase_							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)

          # Now test with buffers included in the offload
          attach_align_device_hook(
              lowercase_   ,       execution_device=lowercase_   ,       offload=lowercase_   ,       weights_map=model.state_dict()   ,       offload_buffers=lowercase_   ,       )

          # Parameters have been offloaded, so on the meta device, buffers included
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('meta'							)							)
          self.assertEqual(model.batchnorm.running_mean.device   ,       torch.device('meta'							)							)

          UpperCAmelCase									=					torch.randn(2   ,       3							)
          UpperCAmelCase									=					model(lowercase_							)
          self.assertEqual(output.device   ,       lowercase_							)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(lowercase_							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.batchnorm.weight.device   ,       torch.device('cpu'							)							)
          self.assertEqual(model.lineara.weight.device   ,       torch.device('cpu'							)							)



