


import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_MAPPING,
        DeiTForImageClassification,
        DeiTForImageClassificationWithTeacher,
        DeiTForMaskedImageModeling,
        DeiTModel,
    )
    from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import DeiTImageProcessor


class     lowercase__  :






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




         """simple docstring"""







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

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






  def 			_UpperCAmelCase  (					self      :					Optional[Any] ):




         """simple docstring"""







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

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

         UpperCAmelCase__										=     self.get_config()

         return config, pixel_values, labels






  def 			_UpperCAmelCase  (					self      :					List[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=_lowercase		,						initializer_range=self.initializer_range		,						encoder_stride=self.encoder_stride		,						)






  def 			_UpperCAmelCase  (					self      :					Union[str, Any]		,						_lowercase      :					str		,						_lowercase      :					int		,						_lowercase      :					Any ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					List[Any]		,						_lowercase      :					Optional[int]		,						_lowercase      :					int		,						_lowercase      :					Tuple ):




         """simple docstring"""







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

         # test greyscale images
         UpperCAmelCase__										=     1
         UpperCAmelCase__										=     DeiTForMaskedImageModeling(_lowercase )
         model.to(_lowercase )
         model.eval()

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






  def 			_UpperCAmelCase  (					self      :					Tuple		,						_lowercase      :					Any		,						_lowercase      :					Optional[int]		,						_lowercase      :					Tuple ):




         """simple docstring"""







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

         # test greyscale images
         UpperCAmelCase__										=     1
         UpperCAmelCase__										=     DeiTForImageClassification(_lowercase )
         model.to(_lowercase )
         model.eval()

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






  def 			_UpperCAmelCase  (					self      :					Optional[Any] ):




         """simple docstring"""







         UpperCAmelCase__										=     self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase__
         )						,			(
             UpperCAmelCase__
         )						,			(
             UpperCAmelCase__
         )						,			
         )										=     config_and_inputs
         UpperCAmelCase__										=     {"pixel_values": pixel_values}
         return config, inputs_dict




@require_torch
class     lowercase__  (					__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							unittest.TestCase		):
  A__=  (
      (
          DeiTModel,
          DeiTForImageClassification,
          DeiTForImageClassificationWithTeacher,
          DeiTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  A__=  (
      {
          'feature-extraction': DeiTModel,
          'image-classification': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  A__=  False
  A__=  False
  A__=  False






  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         self.config_tester.run_common_tests()






  @unittest.skip(reason="DeiT does not use inputs_embeds" )
  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         pass






  def 			_UpperCAmelCase  (					self      :					str ):




         """simple docstring"""







         UpperCAmelCase__						,			UpperCAmelCase__										=     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                UpperCAmelCase__										=     model_class(_lowercase )
                self.assertIsInstance(model.get_input_embeddings()		,						(nn.Module) )
                UpperCAmelCase__										=     model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(_lowercase		,						nn.Linear ) )






  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         UpperCAmelCase__						,			UpperCAmelCase__										=     self.model_tester.prepare_config_and_inputs_for_common()

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

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






  def 			_UpperCAmelCase  (					self      :					str ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					List[str] ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					Tuple ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					List[str]		,						_lowercase      :					str		,						_lowercase      :					Union[str, Any]		,						_lowercase      :					int=False ):




         """simple docstring"""







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

         if return_labels:
                if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
                       del inputs_dict["labels"]

         return inputs_dict






  def 			_UpperCAmelCase  (					self      :					List[str] ):




         """simple docstring"""







         if not self.model_tester.is_training:
                return

         UpperCAmelCase__						,			UpperCAmelCase__										=     self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__										=     True

         for model_class in self.all_model_classes:
                # DeiTForImageClassificationWithTeacher supports inference-only
                if (
                    model_class in get_values(_lowercase )
                    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                ):
                       continue
                UpperCAmelCase__										=     model_class(_lowercase )
                model.to(_lowercase )
                model.train()
                UpperCAmelCase__										=     self._prepare_for_class(_lowercase		,						_lowercase		,						return_labels=_lowercase )
                UpperCAmelCase__										=     model(**_lowercase ).loss
                loss.backward()






  def 			_UpperCAmelCase  (					self      :					Dict ):




         """simple docstring"""







         UpperCAmelCase__						,			UpperCAmelCase__										=     self.model_tester.prepare_config_and_inputs_for_common()
         if not self.model_tester.is_training:
                return

         UpperCAmelCase__										=     False
         UpperCAmelCase__										=     True

         for model_class in self.all_model_classes:
                if model_class in get_values(_lowercase ) or not model_class.supports_gradient_checkpointing:
                       continue
                # DeiTForImageClassificationWithTeacher supports inference-only
                if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
                       continue
                UpperCAmelCase__										=     model_class(_lowercase )
                model.gradient_checkpointing_enable()
                model.to(_lowercase )
                model.train()
                UpperCAmelCase__										=     self._prepare_for_class(_lowercase		,						_lowercase		,						return_labels=_lowercase )
                UpperCAmelCase__										=     model(**_lowercase ).loss
                loss.backward()






  def 			_UpperCAmelCase  (					self      :					Tuple ):




         """simple docstring"""







         UpperCAmelCase__						,			UpperCAmelCase__										=     self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase__										=     [
             {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
             {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
             {"title": "regression", "num_labels": 1, "dtype": torch.float},
         ]

         for model_class in self.all_model_classes:
                if (
                    model_class
                    not in [
                        *get_values(_lowercase ),
                        *get_values(_lowercase ),
                    ]
                    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                ):
                       continue

                for problem_type in problem_types:
                       with self.subTest(msg=F"""Testing {model_class} with {problem_type["title"]}""" ):
                              UpperCAmelCase__										=     problem_type["title"]
                              UpperCAmelCase__										=     problem_type["num_labels"]

                              UpperCAmelCase__										=     model_class(_lowercase )
                              model.to(_lowercase )
                              model.train()

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

                              if problem_type["num_labels"] > 1:
                                     UpperCAmelCase__										=     inputs["labels"].unsqueeze(1 ).repeat(1		,						problem_type["num_labels"] )

                              UpperCAmelCase__										=     inputs["labels"].to(problem_type["dtype"] )

                              # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                              # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                              # they have the same size." which is a symptom something in wrong for the regression problem.
                              # See https://github.com/huggingface/transformers/issues/11780
                              with warnings.catch_warnings(record=_lowercase ) as warning_list:
                                     UpperCAmelCase__										=     model(**_lowercase ).loss
                              for w in warning_list:
                                     if "Using a target size that is different to the input size" in str(w.message ):
                                            raise ValueError(
                                                F"""Something is going wrong in the regression problem: intercepted {w.message}""" )

                              loss.backward()






  @slow
  def 			_UpperCAmelCase  (					self      :					Dict ):




         """simple docstring"""







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





def __UpperCAmelCase       (		)						->					int:



       '''simple docstring'''





       UpperCAmelCase__										=     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
       return image




@require_torch
@require_vision
class     lowercase__  (					unittest.TestCase		):






  @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"""







         UpperCAmelCase__										=     DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224" ).to(
             _lowercase )

         UpperCAmelCase__										=     self.default_image_processor
         UpperCAmelCase__										=     prepare_img()
         UpperCAmelCase__										=     image_processor(images=_lowercase		,						return_tensors="pt" ).to(_lowercase )

         # forward pass
         with torch.no_grad():
                UpperCAmelCase__										=     model(**_lowercase )

         # verify the logits
         UpperCAmelCase__										=     torch.Size((1, 10_00) )
         self.assertEqual(outputs.logits.shape		,						_lowercase )

         UpperCAmelCase__										=     torch.tensor([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1] ).to(_lowercase )

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






  @slow
  @require_accelerate
  @require_torch_gpu
  def 			_UpperCAmelCase  (					self      :					Any ):




         """simple docstring"""







         UpperCAmelCase__										=     DeiTModel.from_pretrained(
             "facebook/deit-base-distilled-patch16-224"		,						torch_dtype=torch.floataa		,						device_map="auto" )
         UpperCAmelCase__										=     self.default_image_processor

         UpperCAmelCase__										=     prepare_img()
         UpperCAmelCase__										=     image_processor(images=_lowercase		,						return_tensors="pt" )
         UpperCAmelCase__										=     inputs.pixel_values.to(_lowercase )

         # forward pass to make sure inference works in fp16
         with torch.no_grad():
                UpperCAmelCase__										=     model(_lowercase )









from .imports import is_tqdm_available


if is_tqdm_available():
    from tqdm.auto import tqdm as _tqdm

from ..state import PartialState





def __UpperCAmelCase       (		__A = True ,						*__A ,						**__A			)						->					Any:



       '''simple docstring'''





       if not is_tqdm_available():
              raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."			)
       UpperCAmelCase__										=     False
       if main_process_only:
              UpperCAmelCase__										=     PartialState().local_process_index == 0
       return _tqdm(*__A ,						**__A ,						disable=__A			)





