




'''simple docstring'''







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						:	List[str]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[Any]=13   ,				lowercase						:	int=30   ,				lowercase						:	Optional[int]=2   ,				lowercase						:	Union[str, Any]=3   ,				lowercase						:	Optional[int]=True   ,				lowercase						:	Dict=True   ,				lowercase						:	Tuple=32   ,				lowercase						:	Union[str, Any]=5   ,				lowercase						:	Tuple=4   ,				lowercase						:	str=37   ,				lowercase						:	Union[str, Any]="gelu"   ,				lowercase						:	Optional[int]=0.1   ,				lowercase						:	Dict=0.1   ,				lowercase						:	List[str]=10   ,				lowercase						:	int=0.02   ,				lowercase						:	Dict=3   ,				lowercase						:	int=None   ,				lowercase						:	str=2   ,				):



           '''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     A      (	self						:	Tuple				):



           '''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     A      (	self						:	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     A      (	self						:	Tuple   ,				lowercase						:	Optional[Any]   ,				lowercase						:	List[Any]   ,				lowercase						:	Optional[Any]				):



           '''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     A      (	self						:	List[Any]   ,				lowercase						:	Dict   ,				lowercase						:	str   ,				lowercase						:	Any				):



           '''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     A      (	self						:	Tuple   ,				lowercase						:	List[Any]   ,				lowercase						:	Any   ,				lowercase						:	List[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     A      (	self						:	Tuple				):



           '''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				:     Dict   				=					(
           {
               """feature-extraction""": DeiTModel,
               """image-classification""": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
           }
           if is_torch_available()
           else {}
       )

       __a				:     Dict   				=					False
       __a				:     Union[str, Any]   				=					False
       __a				:     Tuple   				=					False
       def     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		DeiTModelTester(self				)
           UpperCAmelCase			       =		ConfigTester(self   ,				config_class=lowerCamelCase__   ,				has_text_modality=lowerCamelCase__   ,				hidden_size=37				)
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           self.config_tester.run_common_tests()
       @unittest.skip(reason='''DeiT does not use inputs_embeds'''				)
       def     A      (	self						:	str				):



           '''simple docstring'''

           pass
       def     A      (	self						:	Tuple				):



           '''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     A      (	self						:	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(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     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase__				)
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase__				)
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__				)
       def     A      (	self						:	Optional[Any]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	int   ,				lowercase						:	List[str]=False				):



           '''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     A      (	self						:	List[str]				):



           '''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     A      (	self						:	Union[str, 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(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     A      (	self						:	List[str]				):



           '''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     A      (	self						:	int				):



           '''simple docstring'''

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






def 						snake_case_							():
    UpperCAmelCase			       =		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''     )
    return image







@require_torch
@require_vision
class 			_a  (	unittest.TestCase ):
       @cached_property
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

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



           '''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, 1_000)				)
           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     A      (	self						:	List[str]				):



           '''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 os

import numpy
import onnx



def   a_						(					lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							a.name
				UpperCAmelCase__             =							b.name

				UpperCAmelCase__             =							''
				UpperCAmelCase__             =							''

				UpperCAmelCase__             =							a == b

				UpperCAmelCase__             =							name_a
				UpperCAmelCase__             =							name_b

				return res



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				for i, input_name in enumerate(node_proto.input       ):
								if input_name == name:
												node_proto.input.insert(lowerCamelCase					,			lowerCamelCase       )
												node_proto.input.pop(i + 1       )

				if node_proto.op_type == "If":
								_graph_replace_input_with(node_proto.attribute[0].g					,			lowerCamelCase					,			lowerCamelCase       )
								_graph_replace_input_with(node_proto.attribute[1].g					,			lowerCamelCase					,			lowerCamelCase       )
				if node_proto.op_type == "Loop":
								_graph_replace_input_with(node_proto.attribute[0].g					,			lowerCamelCase					,			lowerCamelCase       )



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				for n in graph_proto.node:
								_node_replace_input_with(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       )



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							list(model.graph.initializer       )
				UpperCAmelCase__             =							list(model_without_ext.graph.initializer       )
				for i, ref_i in ind_to_replace:
								assert inits_with_data[i].name == inits[i].name
								assert inits_with_data[ref_i].name == inits[ref_i].name
								assert i > ref_i

								UpperCAmelCase__             =							inits[i].name
								UpperCAmelCase__             =							inits[ref_i].name

								model_without_ext.graph.initializer.remove(inits[i]       )

								# for n in model.graph.node:
								_graph_replace_input_with(model_without_ext.graph					,			lowerCamelCase					,			lowerCamelCase       )







def   a_						(					lowerCamelCase       ):
				UpperCAmelCase__             =							os.path.dirname(lowerCamelCase       )
				UpperCAmelCase__             =							os.path.basename(lowerCamelCase       )

				UpperCAmelCase__             =							onnx.load(os.path.join(lowerCamelCase					,			lowerCamelCase       )       )

				UpperCAmelCase__             =							list(model.graph.initializer       )

				UpperCAmelCase__             =							set()
				UpperCAmelCase__             =							{}
				UpperCAmelCase__             =							[]

				UpperCAmelCase__             =							0

				for i in range(len(lowerCamelCase       )       ):
								if i in dup_set:
												continue

								for j in range(i + 1					,			len(lowerCamelCase       )       ):
												if j in dup_set:
																continue
												if _is_equal_tensor_proto(inits[i]					,			inits[j]       ):
																dup_set.add(lowerCamelCase       )
																dup_set.add(lowerCamelCase       )

																UpperCAmelCase__             =							inits[j].data_type
																UpperCAmelCase__             =							numpy.prod(inits[j].dims       )
																if dtype == 1:
																				mem_size *= 4
																elif dtype == 6:
																				mem_size *= 4
																elif dtype == 7 or dtype == 1_1:
																				mem_size *= 8
																else:
																				print('unexpected data type: '					,			lowerCamelCase       )
																total_reduced_size += mem_size

																UpperCAmelCase__             =							inits[i].name
																UpperCAmelCase__             =							inits[j].name

																if name_i in dup_map:
																				dup_map[name_i].append(lowerCamelCase       )
																else:
																				UpperCAmelCase__             =							[name_j]
																ind_to_replace.append((j, i)       )

				print('total reduced size: '					,			total_reduced_size / 1_0_2_4 / 1_0_2_4 / 1_0_2_4					,			'GB'       )

				UpperCAmelCase__             =							sorted(lowerCamelCase       )
				_remove_dup_initializers_from_model(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       )

				UpperCAmelCase__             =							'optimized_' + model_file_name
				UpperCAmelCase__             =							os.path.join(lowerCamelCase					,			lowerCamelCase       )
				onnx.save(lowerCamelCase					,			lowerCamelCase       )

				return new_model



