





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			,							lowercase_			,							lowercase_=1_3			,							lowercase_=3_0			,							lowercase_=2			,							lowercase_=3			,							lowercase_=True			,							lowercase_=True			,							lowercase_=3_2			,							lowercase_=5			,							lowercase_=4			,							lowercase_=3_7			,							lowercase_="gelu"			,							lowercase_=0.1			,							lowercase_=0.1			,							lowercase_=1_0			,							lowercase_=0.02			,							lowercase_=3			,							lowercase_=None			,							lowercase_=2			,							)     ->       List[Any]:
       __snake_case         =			parent
       __snake_case         =			batch_size
       __snake_case         =			image_size
       __snake_case         =			patch_size
       __snake_case         =			num_channels
       __snake_case         =			is_training
       __snake_case         =			use_labels
       __snake_case         =			hidden_size
       __snake_case         =			num_hidden_layers
       __snake_case         =			num_attention_heads
       __snake_case         =			intermediate_size
       __snake_case         =			hidden_act
       __snake_case         =			hidden_dropout_prob
       __snake_case         =			attention_probs_dropout_prob
       __snake_case         =			type_sequence_label_size
       __snake_case         =			initializer_range
       __snake_case         =			scope
       __snake_case         =			encoder_stride

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






      def 			_a						(		self)     ->       Tuple:
       __snake_case         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

       __snake_case         =			self.get_config()

       return config, pixel_values, labels






      def 			_a						(		self)     ->       Any:
       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 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_)     ->       Union[str, Any]:
       __snake_case         =			DeiTModel(config=_lowercase)
       model.to(_lowercase)
       model.eval()
       __snake_case         =			model(_lowercase)
       self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.seq_length, self.hidden_size))






      def 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_)     ->       Optional[Any]:
       __snake_case         =			DeiTForMaskedImageModeling(config=_lowercase)
       model.to(_lowercase)
       model.eval()
       __snake_case         =			model(_lowercase)
       self.parent.assertEqual(
           result.reconstruction.shape			,							(self.batch_size, self.num_channels, self.image_size, self.image_size))

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

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






      def 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_)     ->       List[Any]:
       __snake_case         =			self.type_sequence_label_size
       __snake_case         =			DeiTForImageClassification(_lowercase)
       model.to(_lowercase)
       model.eval()
       __snake_case         =			model(_lowercase			,							labels=_lowercase)
       self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.type_sequence_label_size))

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

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







      def 			_a						(		self)     ->       int:
       __snake_case         =			self.prepare_config_and_inputs()
       (
           __snake_case			
       )         =			config_and_inputs
       __snake_case         =			{"pixel_values": pixel_values}
       return config, inputs_dict






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

      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False






      def 			_a						(		self)     ->       int:
       __snake_case         =			DeiTModelTester(self)
       __snake_case         =			ConfigTester(self			,							config_class=_lowercase			,							has_text_modality=_lowercase			,							hidden_size=3_7)






      def 			_a						(		self)     ->       Tuple:
       self.config_tester.run_common_tests()






      @unittest.skip(reason='DeiT does not use inputs_embeds')
      def 			_a						(		self)     ->       Optional[Any]:
       pass






      def 			_a						(		self)     ->       List[Any]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()

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






      def 			_a						(		self)     ->       Dict:
       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()

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

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






      def 			_a						(		self)     ->       str:
       __snake_case         =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowercase)






      def 			_a						(		self)     ->       Union[str, Any]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*_lowercase)






      def 			_a						(		self)     ->       int:
       __snake_case         =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_lowercase)






      def 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_=False)     ->       int:
       __snake_case         =			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 			_a						(		self)     ->       List[Any]:
       if not self.model_tester.is_training:
        return

       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()
       __snake_case         =			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
        __snake_case         =			model_class(_lowercase)
        model.to(_lowercase)
        model.train()
        __snake_case         =			self._prepare_for_class(_lowercase			,							_lowercase			,							return_labels=_lowercase)
        __snake_case         =			model(**_lowercase).loss
        loss.backward()






      def 			_a						(		self)     ->       Optional[int]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()
       if not self.model_tester.is_training:
        return

       __snake_case         =			False
       __snake_case         =			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
        __snake_case         =			model_class(_lowercase)
        model.gradient_checkpointing_enable()
        model.to(_lowercase)
        model.train()
        __snake_case         =			self._prepare_for_class(_lowercase			,							_lowercase			,							return_labels=_lowercase)
        __snake_case         =			model(**_lowercase).loss
        loss.backward()






      def 			_a						(		self)     ->       Optional[Any]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()

       __snake_case         =			[
           {"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']}"):
          __snake_case         =			problem_type["title"]
          __snake_case         =			problem_type["num_labels"]

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

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

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

          __snake_case         =			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:
           __snake_case         =			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 			_a						(		self)     ->       Any:
       for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __snake_case         =			DeiTModel.from_pretrained(_lowercase)
        self.assertIsNotNone(_lowercase)






def  A							(			)					->				Any:
 '''simple docstring'''
 __snake_case         =			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
 return image







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






      @cached_property
      def 			_a						(		self)     ->       int:
       return (
           DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224')
           if is_vision_available()
           else None
       )






      @slow
      def 			_a						(		self)     ->       int:
       __snake_case         =			DeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224').to(
           _lowercase)

       __snake_case         =			self.default_image_processor
       __snake_case         =			prepare_img()
       __snake_case         =			image_processor(images=_lowercase			,							return_tensors='pt').to(_lowercase)

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

       # verify the logits
       __snake_case         =			torch.Size((1, 1_0_0_0))
       self.assertEqual(outputs.logits.shape			,							_lowercase)

       __snake_case         =			torch.tensor([-1.0266, 0.1912, -1.2861]).to(_lowercase)

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







      @slow
      @require_accelerate
      @require_torch_gpu
      def 			_a						(		self)     ->       List[Any]:
       __snake_case         =			DeiTModel.from_pretrained(
           'facebook/deit-base-distilled-patch16-224'			,							torch_dtype=torch.floataa			,							device_map='auto')
       __snake_case         =			self.default_image_processor

       __snake_case         =			prepare_img()
       __snake_case         =			image_processor(images=_lowercase			,							return_tensors='pt')
       __snake_case         =			inputs.pixel_values.to(_lowercase)

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






"""simple docstring"""







from random import randint, random

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							int     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							bool = False     ,		lowerCamelCase_:							bool = False     ,		lowerCamelCase_:							int = 5     ,		):
						"""simple docstring"""






						snake_case   :			str						    = [[-1] * number_of_cells]  # Create a highway without any car
						snake_case   :			str						    = 0
						snake_case   :			Any						    = max(lowerCamelCase_     ,		0 )
						while i < number_of_cells:
												snake_case   :			Optional[int]						    = (
												    randint(0     ,		lowerCamelCase_ ) if random_speed else initial_speed
												)  # Place the cars
												i += (
												    randint(1     ,		max_speed * 2 ) if random_frequency else frequency
												)  # Arbitrary number, may need tuning
						return highway

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							list     ,		lowerCamelCase_:							int ):
						"""simple docstring"""






						snake_case   :			Any						    = 0
						snake_case   :			int						    = highway_now[car_index + 1 :]
						for cell in range(len(lowerCamelCase_ ) ):  # May need a better name for this
												if cells[cell] != -1:  # If the cell is not empty then
																		return distance  # we have the distance we wanted
												distance += 1
						# Here if the car is near the end of the highway
						return distance + get_distance(lowerCamelCase_     ,		-1 )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							list     ,		lowerCamelCase_:							float     ,		lowerCamelCase_:							int ):
						"""simple docstring"""






						snake_case   :			List[str]						    = len(lowerCamelCase_ )
						# Beforce calculations, the highway is empty
						snake_case   :			Dict						    = [-1] * number_of_cells

						for car_index in range(lowerCamelCase_ ):
												if highway_now[car_index] != -1:
																		# Add 1 to the current speed of the car and cap the speed
																		snake_case   :			Tuple						    = min(highway_now[car_index] + 1     ,		lowerCamelCase_ )
																		# Number of empty cell before the next car
																		snake_case   :			int						    = get_distance(lowerCamelCase_     ,		lowerCamelCase_ ) - 1
																		# We can't have the car causing an accident
																		snake_case   :			str						    = min(next_highway[car_index]     ,		lowerCamelCase_ )
																		if random() < probability:
																								# Randomly, a driver will slow down
																								snake_case   :			Optional[Any]						    = max(next_highway[car_index] - 1     ,		0 )
						return next_highway

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							list     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							float     ,		lowerCamelCase_:							int ):
						"""simple docstring"""






						snake_case   :			Optional[Any]						    = len(highway[0] )

						for i in range(lowerCamelCase_ ):
												snake_case   :			Union[str, Any]						    = update(highway[i]     ,		lowerCamelCase_     ,		lowerCamelCase_ )
												snake_case   :			Optional[int]						    = [-1] * number_of_cells

												for car_index in range(lowerCamelCase_ ):
																		snake_case   :			List[Any]						    = next_speeds_calculated[car_index]
																		if speed != -1:
																								# Change the position based on the speed (with % to create the loop)
																								snake_case   :			Optional[Any]						    = (car_index + speed) % number_of_cells
																								# Commit the change of position
																								snake_case   :			Tuple						    = speed
												highway.append(lowerCamelCase_ )

						return highway


if __name__ == "__main__":
			import doctest

			doctest.testmod()



