from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TFDeiTForImageClassification,
					    TFDeiTForImageClassificationWithTeacher,
					    TFDeiTForMaskedImageModeling,
					    TFDeiTModel,
					)
					from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import DeiTImageProcessor



class 					lowerCAmelCase_      :



				"""simple docstring"""

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






										"""simple docstring"""

										snake_case        =     parent
										snake_case        =     batch_size
										snake_case        =     image_size
										snake_case        =     patch_size
										snake_case        =     num_channels
										snake_case        =     is_training
										snake_case        =     use_labels
										snake_case        =     hidden_size
										snake_case        =     num_hidden_layers
										snake_case        =     num_attention_heads
										snake_case        =     intermediate_size
										snake_case        =     hidden_act
										snake_case        =     hidden_dropout_prob
										snake_case        =     attention_probs_dropout_prob
										snake_case        =     type_sequence_label_size
										snake_case        =     initializer_range
										snake_case        =     scope
										snake_case        =     encoder_stride

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

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

										snake_case        =     self.get_config()

										return config, pixel_values, labels

				def 						snake_case	(			self      ):






										"""simple docstring"""

										return DeiTConfig(
										    image_size=self.image_size						,      patch_size=self.patch_size						,      num_channels=self.num_channels						,      hidden_size=self.hidden_size						,      num_hidden_layers=self.num_hidden_layers						,      num_attention_heads=self.num_attention_heads						,      intermediate_size=self.intermediate_size						,      hidden_act=self.hidden_act						,      hidden_dropout_prob=self.hidden_dropout_prob						,      attention_probs_dropout_prob=self.attention_probs_dropout_prob						,      is_decoder=__lowerCAmelCase						,      initializer_range=self.initializer_range						,      encoder_stride=self.encoder_stride						,      )

				def 						snake_case	(			self						,      lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase      ):






										"""simple docstring"""

										snake_case        =     TFDeiTModel(config=__lowerCAmelCase      )
										snake_case        =     model(__lowerCAmelCase      )
										self.parent.assertEqual(result.last_hidden_state.shape						,      (self.batch_size, self.seq_length, self.hidden_size)      )

				def 						snake_case	(			self						,      lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase      ):






										"""simple docstring"""

										snake_case        =     TFDeiTForMaskedImageModeling(config=__lowerCAmelCase      )
										snake_case        =     model(__lowerCAmelCase      )
										self.parent.assertEqual(
										    result.reconstruction.shape						,      (self.batch_size, self.num_channels, self.image_size, self.image_size)      )

										# test greyscale images
										snake_case        =     1
										snake_case        =     TFDeiTForMaskedImageModeling(__lowerCAmelCase      )

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

				def 						snake_case	(			self						,      lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase      ):






										"""simple docstring"""

										snake_case        =     self.type_sequence_label_size
										snake_case        =     TFDeiTForImageClassification(__lowerCAmelCase      )
										snake_case        =     model(__lowerCAmelCase						,      labels=__lowerCAmelCase      )
										self.parent.assertEqual(result.logits.shape						,      (self.batch_size, self.type_sequence_label_size)      )

										# test greyscale images
										snake_case        =     1
										snake_case        =     TFDeiTForImageClassification(__lowerCAmelCase      )

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

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     self.prepare_config_and_inputs()
										snake_case   ,snake_case   ,snake_case        =     config_and_inputs
										snake_case        =     {'pixel_values': pixel_values}
										return config, inputs_dict





@require_tf
class 					lowerCAmelCase_      (     lowerCAmelCase_       ,      lowerCAmelCase_       ,      unittest.TestCase   ):



				"""simple docstring"""
				_lowerCAmelCase       :      Optional[Any]     					=  (
				    (
				        TFDeiTModel,
				        TFDeiTForImageClassification,
				        TFDeiTForImageClassificationWithTeacher,
				        TFDeiTForMaskedImageModeling,
				    )
				    if is_tf_available()
				    else ()
				)
				_lowerCAmelCase       :      Tuple     					=  (
				    {
				        'feature-extraction': TFDeiTModel,
				        'image-classification': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
				    }
				    if is_tf_available()
				    else {}
				)

				_lowerCAmelCase       :      List[Any]     					=  False
				_lowerCAmelCase       :      str     					=  False
				_lowerCAmelCase       :      Any     					=  False
				_lowerCAmelCase       :      Union[str, Any]     					=  False

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     TFDeiTModelTester(self      )
										snake_case        =     ConfigTester(self						,      config_class=__lowerCAmelCase						,      has_text_modality=__lowerCAmelCase						,      hidden_size=37      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										self.config_tester.run_common_tests()

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






										"""simple docstring"""

										pass

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case   ,snake_case        =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																snake_case        =     model_class(__lowerCAmelCase      )
																self.assertIsInstance(model.get_input_embeddings()						,      (tf.keras.layers.Layer)      )
																snake_case        =     model.get_output_embeddings()
																self.assertTrue(x is None or isinstance(__lowerCAmelCase						,      tf.keras.layers.Dense      )      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case   ,snake_case        =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																snake_case        =     model_class(__lowerCAmelCase      )
																snake_case        =     inspect.signature(model.call      )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																snake_case        =     [*signature.parameters.keys()]

																snake_case        =     ['pixel_values']
																self.assertListEqual(arg_names[:1]						,      __lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase      )

				def 						snake_case	(			self						,      lowerCAmelCase						,      lowerCAmelCase						,      lowerCAmelCase=False      ):






										"""simple docstring"""

										snake_case        =     super()._prepare_for_class(__lowerCAmelCase						,      __lowerCAmelCase						,      return_labels=__lowerCAmelCase      )

										if return_labels:
																if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call      ).parameters:
																						del inputs_dict["labels"]

										return inputs_dict

				@slow
				def 						snake_case	(			self      ):






										"""simple docstring"""

										for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																snake_case        =     TFDeiTModel.from_pretrained(__lowerCAmelCase      )
																self.assertIsNotNone(__lowerCAmelCase      )





def 				lowerCAmelCase__      (     )   ->			int:





						"""simple docstring"""






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



@require_tf
@require_vision
class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				@cached_property
				def 						snake_case	(			self      ):






										"""simple docstring"""

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

				@slow
				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224'      )

										snake_case        =     self.default_image_processor
										snake_case        =     prepare_img()
										snake_case        =     image_processor(images=__lowerCAmelCase						,      return_tensors='tf'      )

										# forward pass
										snake_case        =     model(**__lowerCAmelCase      )

										# verify the logits
										snake_case        =     tf.TensorShape((1, 10_00)      )
										self.assertEqual(outputs.logits.shape						,      __lowerCAmelCase      )

										snake_case        =     tf.constant([-1.02_66, 0.19_12, -1.28_61]      )

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



"""simple docstring"""
SCREAMING_SNAKE_CASE__   =					{str(digit): digit**5 for digit in range(10)}





def 				lowerCAmelCase__      (     _UpperCamelCase    :			int   )   ->			int:





						"""simple docstring"""






						return sum(DIGITS_FIFTH_POWER[digit] for digit in str(_UpperCamelCase   )   )





def 				lowerCAmelCase__      (     )   ->			int:





						"""simple docstring"""






						return sum(
						    number
						    for number in range(1_0_0_0	,     1_0_0_0_0_0_0   )
						    if number == digits_fifth_powers_sum(_UpperCamelCase   )   )


if __name__ == "__main__":
					print(solution())


