






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

  '''simple docstring'''







  def __init__(      self		:			Tuple       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			str=13       ,				lowerCamelCase		:			Union[str, Any]=30       ,				lowerCamelCase		:			Optional[Any]=2       ,				lowerCamelCase		:			Optional[int]=3       ,				lowerCamelCase		:			Optional[Any]=True       ,				lowerCamelCase		:			Any=True       ,				lowerCamelCase		:			List[str]=32       ,				lowerCamelCase		:			Optional[int]=5       ,				lowerCamelCase		:			Any=4       ,				lowerCamelCase		:			Optional[int]=37       ,				lowerCamelCase		:			Optional[Any]="gelu"       ,				lowerCamelCase		:			Optional[int]=0.1       ,				lowerCamelCase		:			List[str]=0.1       ,				lowerCamelCase		:			List[str]=10       ,				lowerCamelCase		:			Any=0.02       ,				lowerCamelCase		:			List[Any]=3       ,				lowerCamelCase		:			Dict=None       ,				lowerCamelCase		:			Tuple=2       ,				)						->	Any:


         """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  lowerCamelCase       (      self		:			Any							)						->	int:


         """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  lowerCamelCase       (      self		:			int							)						->	Optional[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=lowerCamelCase       ,				initializer_range=self.initializer_range       ,				encoder_stride=self.encoder_stride       ,				)


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			Tuple							)						->	Tuple:


         """simple docstring"""
         _UpperCAmelCase							=    DeiTModel(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase							)
         self.parent.assertEqual(result.last_hidden_state.shape       ,				(self.batch_size, self.seq_length, self.hidden_size)							)


  def  lowerCamelCase       (      self		:			int       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[Any]							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    DeiTForMaskedImageModeling(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase							)
         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(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()

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


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Union[str, Any]							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.type_sequence_label_size
         _UpperCAmelCase							=    DeiTForImageClassification(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.type_sequence_label_size)							)

         # test greyscale images
         _UpperCAmelCase							=    1
         _UpperCAmelCase							=    DeiTForImageClassification(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()

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





  def  lowerCamelCase       (      self		:			int							)						->	Union[str, 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 		SCREAMING_SNAKE_CASE__	(  UpperCAmelCase					,      UpperCAmelCase					,      unittest.TestCase     ):

  '''simple docstring'''







  _lowerCamelCase    =  (
      (
          DeiTModel,
          DeiTForImageClassification,
          DeiTForImageClassificationWithTeacher,
          DeiTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase    =  (
      {
          '''feature-extraction''': DeiTModel,
          '''image-classification''': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  _lowerCamelCase    =  False
  _lowerCamelCase    =  False
  _lowerCamelCase    =  False


  def  lowerCamelCase       (      self		:			Dict							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    DeiTModelTester(self							)
         _UpperCAmelCase							=    ConfigTester(self       ,				config_class=lowerCamelCase       ,				has_text_modality=lowerCamelCase       ,				hidden_size=37							)


  def  lowerCamelCase       (      self		:			Tuple							)						->	Tuple:


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


  @unittest.skip(reason="""DeiT does not use inputs_embeds"""							)
  def  lowerCamelCase       (      self		:			str							)						->	Union[str, Any]:


         """simple docstring"""
         pass


  def  lowerCamelCase       (      self		:			Dict							)						->	Optional[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(lowerCamelCase							)
                self.assertIsInstance(model.get_input_embeddings()       ,				(nn.Module)							)
                _UpperCAmelCase							=    model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(lowerCamelCase       ,				nn.Linear							)							)


  def  lowerCamelCase       (      self		:			Any							)						->	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(lowerCamelCase							)
                _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]       ,				lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			List[str]							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Union[str, Any]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			Optional[Any]=False							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    super()._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)

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

         return inputs_dict


  def  lowerCamelCase       (      self		:			List[str]							)						->	Dict:


         """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(lowerCamelCase							)
                    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                ):
                       continue
                _UpperCAmelCase							=    model_class(lowerCamelCase							)
                model.to(lowerCamelCase							)
                model.train()
                _UpperCAmelCase							=    self._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)
                _UpperCAmelCase							=    model(**lowerCamelCase							).loss
                loss.backward()


  def  lowerCamelCase       (      self		:			int							)						->	Tuple:


         """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(lowerCamelCase							) or not model_class.supports_gradient_checkpointing:
                       continue
                # DeiTForImageClassificationWithTeacher supports inference-only
                if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
                       continue
                _UpperCAmelCase							=    model_class(lowerCamelCase							)
                model.gradient_checkpointing_enable()
                model.to(lowerCamelCase							)
                model.train()
                _UpperCAmelCase							=    self._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)
                _UpperCAmelCase							=    model(**lowerCamelCase							).loss
                loss.backward()


  def  lowerCamelCase       (      self		:			Dict							)						->	int:


         """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(lowerCamelCase							),
                        *get_values(lowerCamelCase							),
                    ]
                    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(lowerCamelCase							)
                              model.to(lowerCamelCase							)
                              model.train()

                              _UpperCAmelCase							=    self._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)

                              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=lowerCamelCase							) as warning_list:
                                     _UpperCAmelCase							=    model(**lowerCamelCase							).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  lowerCamelCase       (      self		:			Tuple							)						->	int:


         """simple docstring"""
         for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _UpperCAmelCase							=    DeiTModel.from_pretrained(lowerCamelCase							)
                self.assertIsNotNone(lowerCamelCase							)




def       _SCREAMING_SNAKE_CASE (							)				->						Union[str, Any]:
       _UpperCAmelCase							=    Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
       return image





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

  '''simple docstring'''







  @cached_property
  def  lowerCamelCase       (      self		:			Any							)						->	Any:


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


  @slow
  def  lowerCamelCase       (      self		:			str							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""							).to(
             lowerCamelCase							)

         _UpperCAmelCase							=    self.default_image_processor
         _UpperCAmelCase							=    prepare_img()
         _UpperCAmelCase							=    image_processor(images=lowerCamelCase       ,				return_tensors="""pt"""							).to(lowerCamelCase							)

         # forward pass
         with torch.no_grad():
                _UpperCAmelCase							=    model(**lowerCamelCase							)

         # verify the logits
         _UpperCAmelCase							=    torch.Size((1, 1000)							)
         self.assertEqual(outputs.logits.shape       ,				lowerCamelCase							)

         _UpperCAmelCase							=    torch.tensor([-1.0266, 0.1912, -1.2861]							).to(lowerCamelCase							)

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





  @slow
  @require_accelerate
  @require_torch_gpu
  def  lowerCamelCase       (      self		:			Optional[int]							)						->	Tuple:


         """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=lowerCamelCase       ,				return_tensors="""pt"""							)
         _UpperCAmelCase							=    inputs.pixel_values.to(lowerCamelCase							)

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




"""simple docstring"""






import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin



class      A__  (    __SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE    ,						__SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	@register_to_config
	def __init__(	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						float			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						bool = False			,     )							->			Any:

				"""simple docstring"""

				super().__init__()

				__lowerCAmelCase :    str				=			nn.Embedding(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			nn.Embedding(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    str				=			False

				__lowerCAmelCase :    Optional[Any]				=			nn.Dropout(p=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			TaConfig(
				    vocab_size=_SCREAMING_SNAKE_CASE			,     d_model=_SCREAMING_SNAKE_CASE			,     num_heads=_SCREAMING_SNAKE_CASE			,     d_kv=_SCREAMING_SNAKE_CASE			,     d_ff=_SCREAMING_SNAKE_CASE			,     dropout_rate=_SCREAMING_SNAKE_CASE			,     feed_forward_proj=_SCREAMING_SNAKE_CASE			,     is_decoder=_SCREAMING_SNAKE_CASE			,     is_encoder_decoder=_SCREAMING_SNAKE_CASE			,     )

				__lowerCAmelCase :    str				=			nn.ModuleList()
				for lyr_num in range(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase :    int				=			TaBlock(_SCREAMING_SNAKE_CASE)
							self.encoders.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Dict				=			TaLayerNorm(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Any				=			nn.Dropout(p=_SCREAMING_SNAKE_CASE)







	def 				_SCREAMING_SNAKE_CASE  (	self:						Any			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						str)							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			self.token_embedder(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Optional[Any]				=			encoder_input_tokens.shape[1]
				__lowerCAmelCase :    List[Any]				=			torch.arange(_SCREAMING_SNAKE_CASE			,     device=encoder_input_tokens.device)
				x += self.position_encoding(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[Any]				=			self.dropout_pre(_SCREAMING_SNAKE_CASE)

				# inverted the attention mask
				__lowerCAmelCase :    List[Any]				=			encoder_input_tokens.size()
				__lowerCAmelCase :    Any				=			self.get_extended_attention_mask(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)

				for lyr in self.encoders:
							__lowerCAmelCase :    Union[str, Any]				=			lyr(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)[0]
				__lowerCAmelCase :    int				=			self.layer_norm(_SCREAMING_SNAKE_CASE)

				return self.dropout_post(_SCREAMING_SNAKE_CASE), encoder_inputs_mask