
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 __snake_case		:
    def __init__(  self		,snake_case		,snake_case=13		,snake_case=30		,snake_case=2		,snake_case=3		,snake_case=True		,snake_case=True		,snake_case=32		,snake_case=5		,snake_case=4		,snake_case=37		,snake_case="gelu"		,snake_case=0.1		,snake_case=0.1		,snake_case=10		,snake_case=0.02		,snake_case=3		,snake_case=None		,snake_case=2		,):



           '''simple docstring'''







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

           # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
           lowercase							:					List[Any]	     =				(image_size // patch_size) ** 2
           lowercase							:					Tuple	     =				num_patches + 2
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

           lowercase							:					Tuple	     =				None
           if self.use_labels:
                  lowercase							:					str	     =				ids_tensor([self.batch_size]		,self.type_sequence_label_size				)

           lowercase							:					Optional[Any]	     =				self.get_config()

           return config, pixel_values, labels
    def 				_SCREAMING_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=snake_case		,initializer_range=self.initializer_range		,encoder_stride=self.encoder_stride		,)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					List[Any]	     =				DeiTModel(config=snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					int	     =				model(snake_case				)
           self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					Any	     =				DeiTForMaskedImageModeling(config=snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					Tuple	     =				model(snake_case				)
           self.parent.assertEqual(
               result.reconstruction.shape		,(self.batch_size, self.num_channels, self.image_size, self.image_size)				)

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

           lowercase							:					Tuple	     =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
           lowercase							:					List[str]	     =				model(snake_case				)
           self.parent.assertEqual(result.reconstruction.shape		,(self.batch_size, 1, self.image_size, self.image_size)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					List[Any]	     =				self.type_sequence_label_size
           lowercase							:					Any	     =				DeiTForImageClassification(snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					int	     =				model(snake_case		,labels=snake_case				)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.type_sequence_label_size)				)

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

           lowercase							:					Any	     =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
           lowercase							:					Optional[int]	     =				model(snake_case		,labels=snake_case				)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.type_sequence_label_size)				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				self.prepare_config_and_inputs()
           (
               lowercase      
           )							:					int	     =				config_and_inputs
           lowercase							:					Any	     =				{"""pixel_values""": pixel_values}
           return config, inputs_dict






@require_torch
class __snake_case		(				lowerCAmelCase ,     lowerCAmelCase ,     unittest.TestCase					):
    _a		:						Union[str, Any]=    (
        (
            DeiTModel,
            DeiTForImageClassification,
            DeiTForImageClassificationWithTeacher,
            DeiTForMaskedImageModeling,
        )
        if is_torch_available()
        else ()
    )
    _a		:						Tuple=    (
        {
            "feature-extraction": DeiTModel,
            "image-classification": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
        }
        if is_torch_available()
        else {}
    )

    _a		:						Dict=    False
    _a		:						Dict=    False
    _a		:						Union[str, Any]=    False
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				DeiTModelTester(self				)
           lowercase							:					Optional[int]	     =				ConfigTester(self		,config_class=snake_case		,has_text_modality=snake_case		,hidden_size=37				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







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



           '''simple docstring'''







           pass
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Union[str, Any]	     =				self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  lowercase							:					Optional[int]	     =				model_class(snake_case				)
                  self.assertIsInstance(model.get_input_embeddings()		,(nn.Module)				)
                  lowercase							:					Tuple	     =				model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(snake_case		,nn.Linear				)				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				self.model_tester.prepare_config_and_inputs_for_common()

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

                  lowercase							:					Optional[Any]	     =				["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]		,snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					str	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Any	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case=False				):



           '''simple docstring'''







           lowercase							:					int	     =				super()._prepare_for_class(snake_case		,snake_case		,return_labels=snake_case				)

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

           return inputs_dict
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           if not self.model_tester.is_training:
                  return

           lowercase							:					Dict	     =				self.model_tester.prepare_config_and_inputs_for_common()
           lowercase							:					List[Any]	     =				True

           for model_class in self.all_model_classes:
                  # DeiTForImageClassificationWithTeacher supports inference-only
                  if (
                      model_class in get_values(snake_case				)
                      or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                  ):
                         continue
                  lowercase							:					str	     =				model_class(snake_case				)
                  model.to(snake_case				)
                  model.train()
                  lowercase							:					int	     =				self._prepare_for_class(snake_case		,snake_case		,return_labels=snake_case				)
                  lowercase							:					Tuple	     =				model(**snake_case				).loss
                  loss.backward()
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					str	     =				self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
                  return

           lowercase							:					Optional[int]	     =				False
           lowercase							:					Tuple	     =				True

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



           '''simple docstring'''







           lowercase							:					int	     =				self.model_tester.prepare_config_and_inputs_for_common()

           lowercase							:					Union[str, Any]	     =				[
               {"""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(snake_case				),
                          *get_values(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']}"				):
                                lowercase							:					int	     =				problem_type["""title"""]
                                lowercase							:					int	     =				problem_type["""num_labels"""]

                                lowercase							:					Tuple	     =				model_class(snake_case				)
                                model.to(snake_case				)
                                model.train()

                                lowercase							:					int	     =				self._prepare_for_class(snake_case		,snake_case		,return_labels=snake_case				)

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

                                lowercase							:					Dict	     =				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=snake_case				) as warning_list:
                                       lowercase							:					Union[str, Any]	     =				model(**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 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase							:					Any	     =				DeiTModel.from_pretrained(snake_case				)
                  self.assertIsNotNone(snake_case				)







def 						_snake_case(  )			->   str:
       lowercase							:					Optional[int]	     =				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
       return image




@require_torch
@require_vision
class __snake_case		(				unittest.TestCase					):
    @cached_property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return (
               DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""				)
               if is_vision_available()
               else None
           )
    @slow
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					List[Any]	     =				DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""				).to(
               snake_case				)

           lowercase							:					Optional[int]	     =				self.default_image_processor
           lowercase							:					str	     =				prepare_img()
           lowercase							:					Union[str, Any]	     =				image_processor(images=snake_case		,return_tensors="""pt"""				).to(snake_case				)

           # forward pass
           with torch.no_grad():
                  lowercase							:					List[Any]	     =				model(**snake_case				)

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

           lowercase							:					Union[str, Any]	     =				torch.tensor([-1.0_266, 0.1_912, -1.2_861]				).to(snake_case				)

           self.assertTrue(torch.allclose(outputs.logits[0, :3]		,snake_case		,atol=1e-4				)				)
    @slow
    @require_accelerate
    @require_torch_gpu
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Dict	     =				DeiTModel.from_pretrained(
               """facebook/deit-base-distilled-patch16-224"""		,torch_dtype=torch.floataa		,device_map="""auto"""				)
           lowercase							:					Union[str, Any]	     =				self.default_image_processor

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

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



from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


lowercase  : List[str]						   =							logging.get_logger(__name__)

lowercase  : Tuple						   =							{
    """google/umt5-small""": """https://huggingface.co/google/umt5-small/resolve/main/config.json""",
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}

class __snake_case		(				lowerCAmelCase					):
    _a		:						List[str]=    "umt5"
    _a		:						Optional[Any]=    ["past_key_values"]
    def __init__(  self		,snake_case=250112		,snake_case=512		,snake_case=64		,snake_case=1024		,snake_case=8		,snake_case=None		,snake_case=6		,snake_case=32		,snake_case=128		,snake_case=0.1		,snake_case=1e-6		,snake_case=1.0		,snake_case="gated-gelu"		,snake_case=True		,snake_case=True		,snake_case="T5Tokenizer"		,snake_case=True		,snake_case=0		,snake_case=1		,snake_case=0		,**snake_case		,):



           '''simple docstring'''







           super().__init__(
               is_encoder_decoder=snake_case		,tokenizer_class=snake_case		,tie_word_embeddings=snake_case		,pad_token_id=snake_case		,eos_token_id=snake_case		,decoder_start_token_id=snake_case		,**snake_case		,)
           lowercase							:					List[str]	     =				vocab_size
           lowercase							:					Optional[Any]	     =				d_model
           lowercase							:					int	     =				d_kv
           lowercase							:					List[Any]	     =				d_ff
           lowercase							:					Dict	     =				num_layers
           lowercase							:					int	     =				(
               num_decoder_layers if num_decoder_layers is not None else self.num_layers
           )  # default = symmetry
           lowercase							:					List[Any]	     =				num_heads
           lowercase							:					Optional[Any]	     =				relative_attention_num_buckets
           lowercase							:					Dict	     =				relative_attention_max_distance
           lowercase							:					Dict	     =				dropout_rate
           lowercase							:					Any	     =				layer_norm_epsilon
           lowercase							:					Any	     =				initializer_factor
           lowercase							:					Union[str, Any]	     =				feed_forward_proj
           lowercase							:					Optional[Any]	     =				use_cache

           lowercase							:					Dict	     =				self.feed_forward_proj.split("""-"""				)
           lowercase							:					List[Any]	     =				act_info[-1]
           lowercase							:					Any	     =				act_info[0] == """gated"""

           if len(snake_case				) > 1 and act_info[0] != "gated" or len(snake_case				) > 2:
                  raise ValueError(
                      f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
                      """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
                      """'gated-gelu' or 'relu'"""				)

           if feed_forward_proj == "gated-gelu":
                  lowercase							:					int	     =				"""gelu_new"""
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return self.d_model
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return self.num_heads
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return self.num_layers






class __snake_case		(				lowerCAmelCase					):
    @property
    # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					int	     =				{
               """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
               """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
           }
           if self.use_past:
                  lowercase							:					List[Any]	     =				"""past_encoder_sequence + sequence"""
                  lowercase							:					Dict	     =				{0: """batch"""}
                  lowercase							:					Union[str, Any]	     =				{0: """batch""", 1: """past_decoder_sequence + sequence"""}
           else:
                  lowercase							:					Optional[Any]	     =				{0: """batch""", 1: """decoder_sequence"""}
                  lowercase							:					List[Any]	     =				{0: """batch""", 1: """decoder_sequence"""}

           if self.use_past:
                  self.fill_with_past_key_values_(snake_case		,direction="""inputs"""				)

           return common_inputs
    @property
    # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return 13
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return 5e-4
