

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





  '''simple docstring'''



  def __init__(  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_=13   ,SCREAMING_SNAKE_CASE_=30   ,SCREAMING_SNAKE_CASE_=2   ,SCREAMING_SNAKE_CASE_=3   ,SCREAMING_SNAKE_CASE_=True   ,SCREAMING_SNAKE_CASE_=True   ,SCREAMING_SNAKE_CASE_=32   ,SCREAMING_SNAKE_CASE_=5   ,SCREAMING_SNAKE_CASE_=4   ,SCREAMING_SNAKE_CASE_=37   ,SCREAMING_SNAKE_CASE_="gelu"   ,SCREAMING_SNAKE_CASE_=0.1   ,SCREAMING_SNAKE_CASE_=0.1   ,SCREAMING_SNAKE_CASE_=10   ,SCREAMING_SNAKE_CASE_=0.02   ,SCREAMING_SNAKE_CASE_=3   ,SCREAMING_SNAKE_CASE_=None   ,SCREAMING_SNAKE_CASE_=2   ,):

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

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





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		List[str]        =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

      snake_case							:		List[Any]        =		None
      if self.use_labels:
          snake_case							:		Optional[int]        =		ids_tensor([self.batch_size]   ,self.type_sequence_label_size    )

      snake_case							:		List[Any]        =		self.get_config()

      return config, pixel_values, labels





  def    snake_case_       (  self    ):

      '''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=SCREAMING_SNAKE_CASE_   ,initializer_range=self.initializer_range   ,encoder_stride=self.encoder_stride   ,)





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		Dict        =		DeiTModel(config=SCREAMING_SNAKE_CASE_    )
      model.to(SCREAMING_SNAKE_CASE_    )
      model.eval()
      snake_case							:		Optional[int]        =		model(SCREAMING_SNAKE_CASE_    )
      self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		Union[str, Any]        =		DeiTForMaskedImageModeling(config=SCREAMING_SNAKE_CASE_    )
      model.to(SCREAMING_SNAKE_CASE_    )
      model.eval()
      snake_case							:		List[Any]        =		model(SCREAMING_SNAKE_CASE_    )
      self.parent.assertEqual(
          result.reconstruction.shape   ,(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

      # test greyscale images
      snake_case							:		str        =		1
      snake_case							:		Union[str, Any]        =		DeiTForMaskedImageModeling(SCREAMING_SNAKE_CASE_    )
      model.to(SCREAMING_SNAKE_CASE_    )
      model.eval()

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





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_    ):

      '''simple docstring'''
      snake_case							:		Any        =		self.type_sequence_label_size
      snake_case							:		Optional[int]        =		DeiTForImageClassification(SCREAMING_SNAKE_CASE_    )
      model.to(SCREAMING_SNAKE_CASE_    )
      model.eval()
      snake_case							:		Any        =		model(SCREAMING_SNAKE_CASE_   ,labels=SCREAMING_SNAKE_CASE_    )
      self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.type_sequence_label_size)    )

      # test greyscale images
      snake_case							:		Any        =		1
      snake_case							:		List[str]        =		DeiTForImageClassification(SCREAMING_SNAKE_CASE_    )
      model.to(SCREAMING_SNAKE_CASE_    )
      model.eval()

      snake_case							:		Union[str, Any]        =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
      snake_case							:		Tuple        =		model(SCREAMING_SNAKE_CASE_   ,labels=SCREAMING_SNAKE_CASE_    )
      self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.type_sequence_label_size)    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		List[Any]        =		self.prepare_config_and_inputs()
      (
          (
          snake_case
      )      ,		(
          snake_case
      )      ,		(
          snake_case
      )      ,		
      )							:		str        =		config_and_inputs
      snake_case							:		List[Any]        =		{"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_torch
class 						_A     (				snake_case						,    snake_case						,    unittest.TestCase				):





  '''simple docstring'''



  __lowerCamelCase			:      List[str]					=   (
      (
          DeiTModel,
          DeiTForImageClassification,
          DeiTForImageClassificationWithTeacher,
          DeiTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  __lowerCamelCase			:      int					=   (
      {
          '''feature-extraction''': DeiTModel,
          '''image-classification''': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  __lowerCamelCase			:      List[str]					=   False
  __lowerCamelCase			:      Optional[Any]					=   False
  __lowerCamelCase			:      List[Any]					=   False





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Optional[int]        =		DeiTModelTester(self    )
      snake_case							:		Union[str, Any]        =		ConfigTester(self   ,config_class=SCREAMING_SNAKE_CASE_   ,has_text_modality=SCREAMING_SNAKE_CASE_   ,hidden_size=37    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      self.config_tester.run_common_tests()





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

      '''simple docstring'''
      pass





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case      ,		snake_case							:		Union[str, Any]        =		self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          snake_case							:		str        =		model_class(SCREAMING_SNAKE_CASE_    )
          self.assertIsInstance(model.get_input_embeddings()   ,(nn.Module)    )
          snake_case							:		List[Any]        =		model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_   ,nn.Linear    )    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case      ,		snake_case							:		List[Any]        =		self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          snake_case							:		str        =		model_class(SCREAMING_SNAKE_CASE_    )
          snake_case							:		Optional[int]        =		inspect.signature(model.forward    )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          snake_case							:		Any        =		[*signature.parameters.keys()]

          snake_case							:		Optional[int]        =		["""pixel_values"""]
          self.assertListEqual(arg_names[:1]   ,SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Any        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		str        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_image_modeling(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		int        =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_    )





  def    snake_case_       (  self   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_=False    ):

      '''simple docstring'''
      snake_case							:		Any        =		super()._prepare_for_class(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,return_labels=SCREAMING_SNAKE_CASE_    )

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

      return inputs_dict





  def    snake_case_       (  self    ):

      '''simple docstring'''
      if not self.model_tester.is_training:
          return

      snake_case      ,		snake_case							:		Optional[Any]        =		self.model_tester.prepare_config_and_inputs_for_common()
      snake_case							:		str        =		True

      for model_class in self.all_model_classes:
          # DeiTForImageClassificationWithTeacher supports inference-only
          if (
              model_class in get_values(SCREAMING_SNAKE_CASE_    )
              or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
          ):
              continue
          snake_case							:		Tuple        =		model_class(SCREAMING_SNAKE_CASE_    )
          model.to(SCREAMING_SNAKE_CASE_    )
          model.train()
          snake_case							:		Union[str, Any]        =		self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,return_labels=SCREAMING_SNAKE_CASE_    )
          snake_case							:		Optional[Any]        =		model(**SCREAMING_SNAKE_CASE_    ).loss
          loss.backward()





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case      ,		snake_case							:		int        =		self.model_tester.prepare_config_and_inputs_for_common()
      if not self.model_tester.is_training:
          return

      snake_case							:		Tuple        =		False
      snake_case							:		Tuple        =		True

      for model_class in self.all_model_classes:
          if model_class in get_values(SCREAMING_SNAKE_CASE_    ) or not model_class.supports_gradient_checkpointing:
              continue
          # DeiTForImageClassificationWithTeacher supports inference-only
          if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
              continue
          snake_case							:		Optional[Any]        =		model_class(SCREAMING_SNAKE_CASE_    )
          model.gradient_checkpointing_enable()
          model.to(SCREAMING_SNAKE_CASE_    )
          model.train()
          snake_case							:		Union[str, Any]        =		self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,return_labels=SCREAMING_SNAKE_CASE_    )
          snake_case							:		Optional[Any]        =		model(**SCREAMING_SNAKE_CASE_    ).loss
          loss.backward()





  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case      ,		snake_case							:		Tuple        =		self.model_tester.prepare_config_and_inputs_for_common()

      snake_case							:		List[str]        =		[
          {"""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(SCREAMING_SNAKE_CASE_    ),
                  *get_values(SCREAMING_SNAKE_CASE_    ),
              ]
              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							:		Any        =		problem_type["""title"""]
                  snake_case							:		str        =		problem_type["""num_labels"""]

                  snake_case							:		str        =		model_class(SCREAMING_SNAKE_CASE_    )
                  model.to(SCREAMING_SNAKE_CASE_    )
                  model.train()

                  snake_case							:		Tuple        =		self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,SCREAMING_SNAKE_CASE_   ,return_labels=SCREAMING_SNAKE_CASE_    )

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

                  snake_case							:		List[Any]        =		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=SCREAMING_SNAKE_CASE_    ) as warning_list:
                      snake_case							:		Any        =		model(**SCREAMING_SNAKE_CASE_    ).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    snake_case_       (  self    ):

      '''simple docstring'''
      for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          snake_case							:		int        =		DeiTModel.from_pretrained(SCREAMING_SNAKE_CASE_    )
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_    )







def       lowercase					(	)				->			List[Any]:





    '''simple docstring'''




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



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





  '''simple docstring'''



  @cached_property
  def    snake_case_       (  self    ):

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





  @slow
  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Optional[Any]        =		DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    ).to(
          SCREAMING_SNAKE_CASE_    )

      snake_case							:		List[Any]        =		self.default_image_processor
      snake_case							:		List[str]        =		prepare_img()
      snake_case							:		Tuple        =		image_processor(images=SCREAMING_SNAKE_CASE_   ,return_tensors="""pt"""    ).to(SCREAMING_SNAKE_CASE_    )

      # forward pass
      with torch.no_grad():
          snake_case							:		int        =		model(**SCREAMING_SNAKE_CASE_    )

      # verify the logits
      snake_case							:		Dict        =		torch.Size((1, 1000)    )
      self.assertEqual(outputs.logits.shape   ,SCREAMING_SNAKE_CASE_    )

      snake_case							:		Union[str, Any]        =		torch.tensor([-1.02_66, 0.19_12, -1.28_61]    ).to(SCREAMING_SNAKE_CASE_    )

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





  @slow
  @require_accelerate
  @require_torch_gpu
  def    snake_case_       (  self    ):

      '''simple docstring'''
      snake_case							:		Dict        =		DeiTModel.from_pretrained(
          """facebook/deit-base-distilled-patch16-224"""   ,torch_dtype=torch.floataa   ,device_map="""auto"""    )
      snake_case							:		List[Any]        =		self.default_image_processor

      snake_case							:		Tuple        =		prepare_img()
      snake_case							:		List[str]        =		image_processor(images=SCREAMING_SNAKE_CASE_   ,return_tensors="""pt"""    )
      snake_case							:		List[str]        =		inputs.pixel_values.to(SCREAMING_SNAKE_CASE_    )

      # forward pass to make sure inference works in fp16
      with torch.no_grad():
          snake_case							:		Optional[Any]        =		model(SCREAMING_SNAKE_CASE_    )








"""simple docstring"""







def       _lowerCamelCase				(							UpperCAmelCase_ :       int     )	-> int:






	"""simple docstring"""





	if not isinstance(UpperCAmelCase_,	UpperCAmelCase_     ):
		raise ValueError("Input must be an integer"     )
	if input_num <= 0:
		raise ValueError("Input must be positive"     )
	return sum(
	    divisor for divisor in range(1,	input_num // 2 + 1     ) if input_num % divisor == 0     )


if __name__ == "__main__":
			import doctest

			doctest.testmod()




