
import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, 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 ViTMAEForPreTraining, ViTMAEModel
		from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import ViTImageProcessor


class SCREAMING_SNAKE_CASE  :






			"""simple docstring"""




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

				"""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_      		= mask_ratio
				snake_case_      		= scope

				# in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
				# (we add 1 for the [CLS] token)
				snake_case_      		= (image_size // patch_size) ** 2
				snake_case_      		= int(math.ceil((1 - mask_ratio) * (num_patches + 1)							)							)


			def 			__lowerCAmelCase (    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 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				return ViTMAEConfig(
				    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=__UpperCamelCase ,				initializer_range=self.initializer_range ,				mask_ratio=self.mask_ratio ,				)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= ViTMAEModel(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= ViTMAEForPreTraining(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase							)
				snake_case_      		= (self.image_size // self.patch_size) ** 2
				snake_case_      		= self.patch_size**2 * self.num_channels
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, num_patches, expected_num_channels)							)

				# test greyscale images
				snake_case_      		= 1
				snake_case_      		= ViTMAEForPreTraining(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
				snake_case_      		= model(__UpperCamelCase							)
				snake_case_      		= self.patch_size**2
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, num_patches, expected_num_channels)							)







			def 			__lowerCAmelCase (    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_torch
class SCREAMING_SNAKE_CASE  (		__snake_case							,						__snake_case							,						unittest.TestCase							):






			"""simple docstring"""




			__A						  =  (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
			__A						  =  {"""feature-extraction""": ViTMAEModel} if is_torch_available() else {}

			__A						  =  False
			__A						  =  False
			__A						  =  False
			__A						  =  False


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= ViTMAEModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase ,				has_text_modality=__UpperCamelCase ,				hidden_size=37							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


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

				"""simple docstring"""


				pass


			def 			__lowerCAmelCase (    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(__UpperCamelCase							)
					self.assertIsInstance(model.get_input_embeddings() ,				(nn.Module)							)
					snake_case_      		= model.get_output_embeddings()
					self.assertTrue(x is None or isinstance(__UpperCamelCase ,				nn.Linear							)							)


			def 			__lowerCAmelCase (    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(__UpperCamelCase							)
					snake_case_      		= inspect.signature(model.forward							)
					# 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] ,				__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_pretraining(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				np.random.seed(2							)

				snake_case_      		= int((pt_model.config.image_size // pt_model.config.patch_size) ** 2							)
				snake_case_      		= np.random.uniform(size=(self.model_tester.batch_size, num_patches)							)
				snake_case_      		= torch.from_numpy(__UpperCamelCase							)

				# Add `noise` argument.
				# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
				snake_case_      		= pt_noise

				super().check_pt_tf_models(__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							)


			def 			__lowerCAmelCase (    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(__UpperCamelCase							)
					model.to(__UpperCamelCase							)
					model.eval()
					# make random mask reproducible
					torch.manual_seed(2							)
					with torch.no_grad():
						snake_case_      		= model(**self._prepare_for_class(__UpperCamelCase ,				__UpperCamelCase							)							)

					snake_case_      		= outputs[0].cpu().numpy()
					snake_case_      		= 0

					with tempfile.TemporaryDirectory() as tmpdirname:
						model.save_pretrained(__UpperCamelCase							)
						snake_case_      		= model_class.from_pretrained(__UpperCamelCase							)
						model.to(__UpperCamelCase							)
						# make random mask reproducible
						torch.manual_seed(2							)
						with torch.no_grad():
							snake_case_      		= model(**self._prepare_for_class(__UpperCamelCase ,				__UpperCamelCase							)							)

						# Make sure we don't have nans
						snake_case_      		= after_outputs[0].cpu().numpy()
						snake_case_      		= 0
						snake_case_      		= np.amax(np.abs(out_a - out_a							)							)
						self.assertLessEqual(__UpperCamelCase ,				1E-5							)


			@unittest.skip(
			    reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass


			@unittest.skip(
			    reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass


			@unittest.skip(
			    reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass


			@unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass


			@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				pass







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					snake_case_      		= ViTMAEModel.from_pretrained(__UpperCamelCase							)
					self.assertIsNotNone(__UpperCamelCase							)







def a():





	'''simple docstring'''
	snake_case_      		= 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							):

				"""simple docstring"""


				return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'							) if is_vision_available() else None







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				np.random.seed(2							)

				snake_case_      		= ViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'							).to(__UpperCamelCase							)

				snake_case_      		= self.default_image_processor
				snake_case_      		= prepare_img()
				snake_case_      		= image_processor(images=__UpperCamelCase ,				return_tensors='pt'							).to(__UpperCamelCase							)

				# prepare a noise vector that will be also used for testing the TF model
				# (this way we can ensure that the PT and TF models operate on the same inputs)
				snake_case_      		= ViTMAEConfig()
				snake_case_      		= int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2							)
				snake_case_      		= np.random.uniform(size=(1, num_patches)							)

				# forward pass
				with torch.no_grad():
					snake_case_      		= model(**__UpperCamelCase ,				noise=torch.from_numpy(__UpperCamelCase							).to(device=__UpperCamelCase							)							)

				# verify the logits
				snake_case_      		= torch.Size((1, 1_96, 7_68)							)
				self.assertEqual(outputs.logits.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]							)

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



from ..utils import DummyObject, requires_backends


class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)
