






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 					SCREAMING_SNAKE_CASE_  :
       """simple docstring"""


       def __init__(				self	:Any,      snake_case	:Tuple,      snake_case	:int=13,      snake_case	:Dict=30,      snake_case	:Any=2,      snake_case	:Any=3,      snake_case	:Any=True,      snake_case	:List[str]=True,      snake_case	:str=32,      snake_case	:Optional[int]=2,      snake_case	:Dict=4,      snake_case	:Dict=37,      snake_case	:List[str]="gelu",      snake_case	:Union[str, Any]=0.1,      snake_case	:Any=0.1,      snake_case	:Union[str, Any]=10,      snake_case	:Any=0.0_2,      snake_case	:Dict=3,      snake_case	:int=None,      snake_case	:Union[str, Any]=2,      ):



          """simple docstring"""

          _lowercase              =parent
          _lowercase              =batch_size
          _lowercase              =image_size
          _lowercase              =patch_size
          _lowercase              =num_channels
          _lowercase              =is_training
          _lowercase              =use_labels
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob
          _lowercase              =type_sequence_label_size
          _lowercase              =initializer_range
          _lowercase              =scope
          _lowercase              =encoder_stride

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





       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase              =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

          _lowercase              =None
          if self.use_labels:
             _lowercase              =ids_tensor([self.batch_size],      self.type_sequence_label_size)

          _lowercase              =self.get_config()

          return config, pixel_values, labels





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """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=UpperCAmelCase__,      initializer_range=self.initializer_range,      encoder_stride=self.encoder_stride,      )





       def 				UpperCamelCase__						(				self	:List[str],      snake_case	:Dict,      snake_case	:List[Any],      snake_case	:Any):



          """simple docstring"""

          _lowercase              =TFDeiTModel(config=UpperCAmelCase__)
          _lowercase              =model(UpperCAmelCase__)
          self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))





       def 				UpperCamelCase__						(				self	:Dict,      snake_case	:Tuple,      snake_case	:str,      snake_case	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =TFDeiTForMaskedImageModeling(config=UpperCAmelCase__)
          _lowercase              =model(UpperCAmelCase__)
          self.parent.assertEqual(
              result.reconstruction.shape,      (self.batch_size, self.num_channels, self.image_size, self.image_size))

          # test greyscale images
          _lowercase              =1
          _lowercase              =TFDeiTForMaskedImageModeling(UpperCAmelCase__)

          _lowercase              =floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
          _lowercase              =model(UpperCAmelCase__)
          self.parent.assertEqual(result.reconstruction.shape,      (self.batch_size, 1, self.image_size, self.image_size))





       def 				UpperCamelCase__						(				self	:Dict,      snake_case	:Union[str, Any],      snake_case	:Any,      snake_case	:List[str]):



          """simple docstring"""

          _lowercase              =self.type_sequence_label_size
          _lowercase              =TFDeiTForImageClassification(UpperCAmelCase__)
          _lowercase              =model(UpperCAmelCase__,      labels=UpperCAmelCase__)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.type_sequence_label_size))

          # test greyscale images
          _lowercase              =1
          _lowercase              =TFDeiTForImageClassification(UpperCAmelCase__)

          _lowercase              =floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
          _lowercase              =model(UpperCAmelCase__,      labels=UpperCAmelCase__)
          self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.type_sequence_label_size))





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =self.prepare_config_and_inputs()
          _lowercase              =config_and_inputs
          _lowercase              ={'''pixel_values''': pixel_values}
          return config, inputs_dict






@require_tf
class 					SCREAMING_SNAKE_CASE_  (       lowercase__  ,						lowercase__  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		Optional[int]					=(
           (
               TFDeiTModel,
               TFDeiTForImageClassification,
               TFDeiTForImageClassificationWithTeacher,
               TFDeiTForMaskedImageModeling,
           )
           if is_tf_available()
           else ()
       )
       __lowerCAmelCase    :		Union[str, Any]					=(
           {
               '''feature-extraction''': TFDeiTModel,
               '''image-classification''': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
           }
           if is_tf_available()
           else {}
       )

       __lowerCAmelCase    :		Optional[int]					=False
       __lowerCAmelCase    :		Dict					=False
       __lowerCAmelCase    :		Union[str, Any]					=False
       __lowerCAmelCase    :		int					=False





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =TFDeiTModelTester(self)
          _lowercase              =ConfigTester(self,      config_class=UpperCAmelCase__,      has_text_modality=UpperCAmelCase__,      hidden_size=37)





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          self.config_tester.run_common_tests()





       @unittest.skip(reason='DeiT does not use inputs_embeds')
       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          pass





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _lowercase              =model_class(UpperCAmelCase__)
             self.assertIsInstance(model.get_input_embeddings(),      (tf.keras.layers.Layer))
             _lowercase              =model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(UpperCAmelCase__,      tf.keras.layers.Dense))





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             _lowercase              =model_class(UpperCAmelCase__)
             _lowercase              =inspect.signature(model.call)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             _lowercase              =[*signature.parameters.keys()]

             _lowercase              =['''pixel_values''']
             self.assertListEqual(arg_names[:1],      UpperCAmelCase__)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*UpperCAmelCase__)





       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase__)





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__)





       def 				UpperCamelCase__						(				self	:Dict,      snake_case	:List[Any],      snake_case	:List[str],      snake_case	:int=False):



          """simple docstring"""

          _lowercase              =super()._prepare_for_class(UpperCAmelCase__,      UpperCAmelCase__,      return_labels=UpperCAmelCase__)

          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 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _lowercase              =TFDeiTModel.from_pretrained(UpperCAmelCase__)
             self.assertIsNotNone(UpperCAmelCase__)





def        _snake_case   ()					->	Optional[Any]:
   _lowercase              =Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
   return image






@require_tf
@require_vision
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       @cached_property
       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

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





       @slow
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          _lowercase              =TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224')

          _lowercase              =self.default_image_processor
          _lowercase              =prepare_img()
          _lowercase              =image_processor(images=UpperCAmelCase__,      return_tensors='tf')

          # forward pass
          _lowercase              =model(**UpperCAmelCase__)

          # verify the logits
          _lowercase              =tf.TensorShape((1, 1000))
          self.assertEqual(outputs.logits.shape,      UpperCAmelCase__)

          _lowercase              =tf.constant([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1])

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





'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import _LazyModule


UpperCamelCase_    						=       {"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
							from .tokenization_byta import ByTaTokenizer
else:
							import sys

							UpperCamelCase_    						=       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





