


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 :
		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     ,       )					->			Optional[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       UpperCAmelCase__							(     self		)					->			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       UpperCAmelCase__							(     self		)					->			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=_SCREAMING_SNAKE_CASE     ,       initializer_range=self.initializer_range     ,       encoder_stride=self.encoder_stride     ,       )
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     DeiTModel(config=_SCREAMING_SNAKE_CASE		)
									model.to(_SCREAMING_SNAKE_CASE		)
									model.eval()
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, self.seq_length, self.hidden_size)		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     DeiTForMaskedImageModeling(config=_SCREAMING_SNAKE_CASE		)
									model.to(_SCREAMING_SNAKE_CASE		)
									model.eval()
									_UpperCAmelCase											=     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
									_UpperCAmelCase											=     1
									_UpperCAmelCase											=     DeiTForMaskedImageModeling(_SCREAMING_SNAKE_CASE		)
									model.to(_SCREAMING_SNAKE_CASE		)
									model.eval()

									_UpperCAmelCase											=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									self.parent.assertEqual(result.reconstruction.shape     ,       (self.batch_size, 1, self.image_size, self.image_size)		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     self.type_sequence_label_size
									_UpperCAmelCase											=     DeiTForImageClassification(_SCREAMING_SNAKE_CASE		)
									model.to(_SCREAMING_SNAKE_CASE		)
									model.eval()
									_UpperCAmelCase											=     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
									_UpperCAmelCase											=     1
									_UpperCAmelCase											=     DeiTForImageClassification(_SCREAMING_SNAKE_CASE		)
									model.to(_SCREAMING_SNAKE_CASE		)
									model.eval()

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






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




									"""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 						__a (							UpperCAmelCase       , UpperCAmelCase       , unittest.TestCase     ):
		_a      :    Dict           =      (
		    (
		        DeiTModel,
		        DeiTForImageClassification,
		        DeiTForImageClassificationWithTeacher,
		        DeiTForMaskedImageModeling,
		    )
		    if is_torch_available()
		    else ()
		)
		_a      :    Optional[int]           =      (
		    {
		        'feature-extraction': DeiTModel,
		        'image-classification': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
		    }
		    if is_torch_available()
		    else {}
		)

		_a      :    Any           =      False
		_a      :    Dict           =      False
		_a      :    List[str]           =      False
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     DeiTModelTester(self		)
									_UpperCAmelCase											=     ConfigTester(self     ,       config_class=_SCREAMING_SNAKE_CASE     ,       has_text_modality=_SCREAMING_SNAKE_CASE     ,       hidden_size=37		)
		def       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									self.config_tester.run_common_tests()
		@unittest.skip(reason='DeiT does not use inputs_embeds'		)
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									pass
		def       UpperCAmelCase__							(     self		)					->			List[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(_SCREAMING_SNAKE_CASE		)
																self.assertIsInstance(model.get_input_embeddings()     ,       (nn.Module)		)
																_UpperCAmelCase											=     model.get_output_embeddings()
																self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE     ,       nn.Linear		)		)
		def       UpperCAmelCase__							(     self		)					->			List[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(_SCREAMING_SNAKE_CASE		)
																_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]     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_image_modeling(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=False		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     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       UpperCAmelCase__							(     self		)					->			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(_SCREAMING_SNAKE_CASE		)
																    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
																):
																							continue
																_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)
																model.to(_SCREAMING_SNAKE_CASE		)
																model.train()
																_UpperCAmelCase											=     self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       return_labels=_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     model(**_SCREAMING_SNAKE_CASE		).loss
																loss.backward()
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""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(_SCREAMING_SNAKE_CASE		) or not model_class.supports_gradient_checkpointing:
																							continue
																# DeiTForImageClassificationWithTeacher supports inference-only
																if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
																							continue
																_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)
																model.gradient_checkpointing_enable()
																model.to(_SCREAMING_SNAKE_CASE		)
																model.train()
																_UpperCAmelCase											=     self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       return_labels=_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase											=     model(**_SCREAMING_SNAKE_CASE		).loss
																loss.backward()
		def       UpperCAmelCase__							(     self		)					->			Optional[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(_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"]}'''		):
																														_UpperCAmelCase											=     problem_type['title']
																														_UpperCAmelCase											=     problem_type['num_labels']

																														_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)
																														model.to(_SCREAMING_SNAKE_CASE		)
																														model.train()

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

																														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=_SCREAMING_SNAKE_CASE		) as warning_list:
																																					_UpperCAmelCase											=     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       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





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



def 			lowerCAmelCase__					(     )    ->						int:


							'''simple docstring'''





							_UpperCAmelCase											=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'			)
							return image







@require_torch
@require_vision
class 						__a (							unittest.TestCase     ):
		@cached_property
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





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




									"""simple docstring"""





									_UpperCAmelCase											=     DeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224'		).to(
									    _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     self.default_image_processor
									_UpperCAmelCase											=     prepare_img()
									_UpperCAmelCase											=     image_processor(images=_SCREAMING_SNAKE_CASE     ,       return_tensors='pt'		).to(_SCREAMING_SNAKE_CASE		)

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

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

									_UpperCAmelCase											=     torch.tensor([-1.0266, 0.1912, -1.2861]		).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       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=_SCREAMING_SNAKE_CASE     ,       return_tensors='pt'		)
									_UpperCAmelCase											=     inputs.pixel_values.to(_SCREAMING_SNAKE_CASE		)

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




from .imports import is_rich_available


if is_rich_available():
		from rich.traceback import install

		install(show_locals=False)

else:
		raise ModuleNotFoundError('''To use the rich extension, install rich with `pip install rich`''')
