

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




  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase=13							,   lowerCAmelCase=30							,   lowerCAmelCase=2							,   lowerCAmelCase=3							,   lowerCAmelCase=True							,   lowerCAmelCase=True							,   lowerCAmelCase=32							,   lowerCAmelCase=2							,   lowerCAmelCase=4							,   lowerCAmelCase=37							,   lowerCAmelCase="gelu"							,   lowerCAmelCase=0.1							,   lowerCAmelCase=0.1							,   lowerCAmelCase=10							,   lowerCAmelCase=0.02							,   lowerCAmelCase=3							,   lowerCAmelCase=None							,   lowerCAmelCase=2							,   )  -> List[Any]:
        '''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   A__     (				self					)  -> Optional[int]:
        '''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   A__     (				self					)  -> str:
        '''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=lowerCAmelCase							,   initializer_range=self.initializer_range							,   encoder_stride=self.encoder_stride							,   )




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)  -> List[str]:
        '''simple docstring'''






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




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)  -> List[Any]:
        '''simple docstring'''






        _lowercase	       =TFDeiTForMaskedImageModeling(config=lowerCAmelCase					)
        _lowercase	       =model(lowerCAmelCase					)
        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(lowerCAmelCase					)

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




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Any:
        '''simple docstring'''






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

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

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




  def   A__     (				self					)  -> Any:
        '''simple docstring'''






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





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

  _a									=  False
  _a									=  False
  _a									=  False
  _a									=  False




  def   A__     (				self					)  -> Any:
        '''simple docstring'''






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




  def   A__     (				self					)  -> Union[str, Any]:
        '''simple docstring'''






        self.config_tester.run_common_tests()




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






        pass




  def   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()

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




  def   A__     (				self					)  -> List[Any]:
        '''simple docstring'''






        _lowercase       ,       _lowercase	       =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              _lowercase	       =model_class(lowerCAmelCase					)
              _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]							,   lowerCAmelCase					)




  def   A__     (				self					)  -> str:
        '''simple docstring'''






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




  def   A__     (				self					)  -> Optional[Any]:
        '''simple docstring'''






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




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






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




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase=False					)  -> Union[str, Any]:
        '''simple docstring'''






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

        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   A__     (				self					)  -> Tuple:
        '''simple docstring'''






        for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              _lowercase	       =TFDeiTModel.from_pretrained(lowerCAmelCase					)
              self.assertIsNotNone(lowerCAmelCase					)
def      a      (			)    ->      str:





      """simple docstring"""



      _lowercase	       =Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
      return image







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




  @cached_property
  def   A__     (				self					)  -> Any:
        '''simple docstring'''






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




  @slow
  def   A__     (				self					)  -> Any:
        '''simple docstring'''






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

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

        # forward pass
        _lowercase	       =model(**lowerCAmelCase					)

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

        _lowercase	       =tf.constant([-1.0266, 0.1912, -1.2861]					)

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





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
  import torch





class  __lowerCAmelCase     (					TensorFormatter[Mapping, """torch.Tensor""", Mapping]							):




  def __init__(				self							,   lowerCAmelCase=None							,   **lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        super().__init__(features=lowerCAmelCase					)
        _lowercase	       =torch_tensor_kwargs
        import torch  # noqa import torch at initialization




  def   A__     (				self							,   lowerCAmelCase					)  -> int:
        '''simple docstring'''






        import torch

        if isinstance(lowerCAmelCase							,   lowerCAmelCase					) and column:
              if all(
                  isinstance(lowerCAmelCase							,   torch.Tensor					) and x.shape == column[0].shape and x.dtype == column[0].dtype
                  for x in column					):
                    return torch.stack(lowerCAmelCase					)
        return column




  def   A__     (				self							,   lowerCAmelCase					)  -> str:
        '''simple docstring'''






        import torch

        if isinstance(lowerCAmelCase							,   (str, bytes, type(lowerCAmelCase					))					):
              return value
        elif isinstance(lowerCAmelCase							,   (np.character, np.ndarray)					) and np.issubdtype(value.dtype							,   np.character					):
              return value.tolist()

        _lowercase	       ={}

        if isinstance(lowerCAmelCase							,   (np.number, np.ndarray)					) and np.issubdtype(value.dtype							,   np.integer					):
              _lowercase	       ={'dtype': torch.intaa}
        elif isinstance(lowerCAmelCase							,   (np.number, np.ndarray)					) and np.issubdtype(value.dtype							,   np.floating					):
              _lowercase	       ={'dtype': torch.floataa}
        elif config.PIL_AVAILABLE and "PIL" in sys.modules:
              import PIL.Image

              if isinstance(lowerCAmelCase							,   PIL.Image.Image					):
                    _lowercase	       =np.asarray(lowerCAmelCase					)
        return torch.tensor(lowerCAmelCase							,   **{**default_dtype, **self.torch_tensor_kwargs}					)




  def   A__     (				self							,   lowerCAmelCase					)  -> str:
        '''simple docstring'''






        import torch

        # support for torch, tf, jax etc.
        if hasattr(lowerCAmelCase							,   '__array__'					) and not isinstance(lowerCAmelCase							,   torch.Tensor					):
              _lowercase	       =data_struct.__array__()
        # support for nested types like struct of list of struct
        if isinstance(lowerCAmelCase							,   np.ndarray					):
              if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
                    return self._consolidate([self.recursive_tensorize(lowerCAmelCase					) for substruct in data_struct]					)
        elif isinstance(lowerCAmelCase							,   (list, tuple)					):
              return self._consolidate([self.recursive_tensorize(lowerCAmelCase					) for substruct in data_struct]					)
        return self._tensorize(lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase					)  -> Tuple:
        '''simple docstring'''






        return map_nested(self._recursive_tensorize							,   lowerCAmelCase							,   map_list=lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase					)  -> Mapping:
        '''simple docstring'''






        _lowercase	       =self.numpy_arrow_extractor().extract_row(lowerCAmelCase					)
        _lowercase	       =self.python_features_decoder.decode_row(lowerCAmelCase					)
        return self.recursive_tensorize(lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase					)  -> "torch.Tensor":
        '''simple docstring'''






        _lowercase	       =self.numpy_arrow_extractor().extract_column(lowerCAmelCase					)
        _lowercase	       =self.python_features_decoder.decode_column(lowerCAmelCase							,   pa_table.column_names[0]					)
        _lowercase	       =self.recursive_tensorize(lowerCAmelCase					)
        _lowercase	       =self._consolidate(lowerCAmelCase					)
        return column




  def   A__     (				self							,   lowerCAmelCase					)  -> Mapping:
        '''simple docstring'''






        _lowercase	       =self.numpy_arrow_extractor().extract_batch(lowerCAmelCase					)
        _lowercase	       =self.python_features_decoder.decode_batch(lowerCAmelCase					)
        _lowercase	       =self.recursive_tensorize(lowerCAmelCase					)
        for column_name in batch:
              _lowercase	       =self._consolidate(batch[column_name]					)
        return batch


