






'''simple docstring'''






from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TFDeiTForImageClassification,
          TFDeiTForImageClassificationWithTeacher,
          TFDeiTForMaskedImageModeling,
          TFDeiTModel,
      )
      from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import DeiTImageProcessor




class a_    :







   def __init__(self					,   __a					,   __a=1_3					,   __a=3_0					,   __a=2					,   __a=3					,   __a=True					,   __a=True					,   __a=3_2					,   __a=2					,   __a=4					,   __a=3_7					,   __a="gelu"					,   __a=0.1					,   __a=0.1					,   __a=1_0					,   __a=0.02					,   __a=3					,   __a=None					,   __a=2					,   )			->      int:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= parent
      __snake_case				:      Tuple     							= batch_size
      __snake_case				:      Union[str, Any]     							= image_size
      __snake_case				:      List[str]     							= patch_size
      __snake_case				:      Tuple     							= num_channels
      __snake_case				:      List[str]     							= is_training
      __snake_case				:      Tuple     							= use_labels
      __snake_case				:      Dict     							= hidden_size
      __snake_case				:      Dict     							= num_hidden_layers
      __snake_case				:      int     							= num_attention_heads
      __snake_case				:      Dict     							= intermediate_size
      __snake_case				:      Any     							= hidden_act
      __snake_case				:      Optional[int]     							= hidden_dropout_prob
      __snake_case				:      Optional[Any]     							= attention_probs_dropout_prob
      __snake_case				:      str     							= type_sequence_label_size
      __snake_case				:      Union[str, Any]     							= initializer_range
      __snake_case				:      Optional[int]     							= scope
      __snake_case				:      List[Any]     							= encoder_stride

      # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
      __snake_case				:      Optional[Any]     							= (image_size // patch_size) ** 2
      __snake_case				:      int     							= num_patches + 2







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      __snake_case				:      Any     							= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

      __snake_case				:      str     							= None
      if self.use_labels:
         __snake_case				:      Tuple     							= ids_tensor([self.batch_size]					,   self.type_sequence_label_size)

      __snake_case				:      List[str]     							= self.get_config()

      return config, pixel_values, labels







   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      return DeiTConfig(
          image_size=self.image_size					,   patch_size=self.patch_size					,   num_channels=self.num_channels					,   hidden_size=self.hidden_size					,   num_hidden_layers=self.num_hidden_layers					,   num_attention_heads=self.num_attention_heads					,   intermediate_size=self.intermediate_size					,   hidden_act=self.hidden_act					,   hidden_dropout_prob=self.hidden_dropout_prob					,   attention_probs_dropout_prob=self.attention_probs_dropout_prob					,   is_decoder=__a					,   initializer_range=self.initializer_range					,   encoder_stride=self.encoder_stride					,   )







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a)			->      List[Any]:
      """simple docstring"""






      __snake_case				:      Any     							= TFDeiTModel(config=__a)
      __snake_case				:      Tuple     							= model(__a)
      self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size))







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a)			->      int:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= TFDeiTForMaskedImageModeling(config=__a)
      __snake_case				:      int     							= model(__a)
      self.parent.assertEqual(
          result.reconstruction.shape					,   (self.batch_size, self.num_channels, self.image_size, self.image_size))

      # test greyscale images
      __snake_case				:      Optional[int]     							= 1
      __snake_case				:      Any     							= TFDeiTForMaskedImageModeling(__a)

      __snake_case				:      Optional[int]     							= floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
      __snake_case				:      Any     							= model(__a)
      self.parent.assertEqual(result.reconstruction.shape					,   (self.batch_size, 1, self.image_size, self.image_size))







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a)			->      List[str]:
      """simple docstring"""






      __snake_case				:      Any     							= self.type_sequence_label_size
      __snake_case				:      Dict     							= TFDeiTForImageClassification(__a)
      __snake_case				:      Union[str, Any]     							= model(__a					,   labels=__a)
      self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.type_sequence_label_size))

      # test greyscale images
      __snake_case				:      Dict     							= 1
      __snake_case				:      List[Any]     							= TFDeiTForImageClassification(__a)

      __snake_case				:      int     							= floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
      __snake_case				:      Dict     							= model(__a					,   labels=__a)
      self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.type_sequence_label_size))







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      __snake_case				:      List[Any]     							= self.prepare_config_and_inputs()
      __snake_case				:      Tuple     							= config_and_inputs
      __snake_case				:      Union[str, Any]     							= {'pixel_values': pixel_values}
      return config, inputs_dict




@require_tf
class a_    (    UpperCamelCase_		,      UpperCamelCase_		,      unittest.TestCase							):
   _snake_case        =							(
       (
           TFDeiTModel,
           TFDeiTForImageClassification,
           TFDeiTForImageClassificationWithTeacher,
           TFDeiTForMaskedImageModeling,
       )
       if is_tf_available()
       else ()
   )
   _snake_case        =							(
       {
           """feature-extraction""": TFDeiTModel,
           """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
       }
       if is_tf_available()
       else {}
   )

   _snake_case        =							False
   _snake_case        =							False
   _snake_case        =							False
   _snake_case        =							False







   def     SCREAMING_SNAKE_CASE__					(self)			->      Tuple:
      """simple docstring"""






      __snake_case				:      Any     							= TFDeiTModelTester(self)
      __snake_case				:      Tuple     							= ConfigTester(self					,   config_class=__a					,   has_text_modality=__a					,   hidden_size=3_7)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      self.config_tester.run_common_tests()







   @unittest.skip(reason='DeiT does not use inputs_embeds')
   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[Any]:
      """simple docstring"""






      __snake_case				:      List[str]     							= self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         __snake_case				:      List[Any]     							= model_class(__a)
         self.assertIsInstance(model.get_input_embeddings()					,   (tf.keras.layers.Layer))
         __snake_case				:      str     							= model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(__a					,   tf.keras.layers.Dense))







   def     SCREAMING_SNAKE_CASE__					(self)			->      Tuple:
      """simple docstring"""






      __snake_case				:      str     							= self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         __snake_case				:      List[str]     							= model_class(__a)
         __snake_case				:      str     							= inspect.signature(model.call)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         __snake_case				:      Union[str, Any]     							= [*signature.parameters.keys()]

         __snake_case				:      str     							= ['pixel_values']
         self.assertListEqual(arg_names[:1]					,   __a)







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      __snake_case				:      int     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__a)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      __snake_case				:      Any     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_image_modeling(*__a)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      Optional[int]     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__a)







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a=False)			->      List[str]:
      """simple docstring"""






      __snake_case				:      List[str]     							= super()._prepare_for_class(__a					,   __a					,   return_labels=__a)

      if return_labels:
         if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call).parameters:
            del inputs_dict["labels"]

      return inputs_dict







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __snake_case				:      List[str]     							= TFDeiTModel.from_pretrained(__a)
         self.assertIsNotNone(__a)




def        _SCREAMING_SNAKE_CASE  (				)    -> Any:







   """simple docstring"""
   __snake_case				:      Optional[Any]     							= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
   return image





@require_tf
@require_vision
class a_    (    unittest.TestCase							):







   @cached_property
   def     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






      return (
          DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224')
          if is_vision_available()
          else None
      )







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      str     							= TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224')

      __snake_case				:      List[str]     							= self.default_image_processor
      __snake_case				:      Any     							= prepare_img()
      __snake_case				:      List[str]     							= image_processor(images=__a					,   return_tensors='tf')

      # forward pass
      __snake_case				:      Tuple     							= model(**__a)

      # verify the logits
      __snake_case				:      List[Any]     							= tf.TensorShape((1, 1_0_0_0))
      self.assertEqual(outputs.logits.shape					,   __a)

      __snake_case				:      List[str]     							= tf.constant([-1.0_266, 0.1_912, -1.2_861])

      self.assertTrue(np.allclose(outputs.logits[0, :3]					,   __a					,   atol=1E-4))







'''simple docstring'''






from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
      import accelerate




def        _SCREAMING_SNAKE_CASE  (				A  :       Optional[Any]  )    -> int:







   """simple docstring"""
   if not is_accelerate_available():
      return method
   __snake_case				:      Optional[Any]     							= version.parse(accelerate.__version__  ).base_version
   if version.parse(A  ) < version.parse('0.17.0'  ):
      return method

   def wrapper(self  :       Optional[Any]  ,		*A  :       Optional[Any]  ,		**A  :       Optional[int]  ):
      if hasattr(self  ,		'_hf_hook'  ) and hasattr(self._hf_hook  ,		'pre_forward'  ):
         self._hf_hook.pre_forward(self  )
      return method(self  ,		*A  ,		**A  )

   return wrapper