





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 						_A:


				"""simple docstring"""


				def __init__(			self     ,					_A     ,					_A=13     ,					_A=30     ,					_A=2     ,					_A=3     ,					_A=True     ,					_A=True     ,					_A=32     ,					_A=5     ,					_A=4     ,					_A=37     ,					_A="gelu"     ,					_A=0.1     ,					_A=0.1     ,					_A=10     ,					_A=0.0_2     ,					_A=3     ,					_A=0.6     ,					_A=None     ,					):
									__A	:						str       =				parent
									__A	:						Union[str, Any]       =				batch_size
									__A	:						Dict       =				image_size
									__A	:						Optional[Any]       =				patch_size
									__A	:						List[Any]       =				num_channels
									__A	:						Union[str, Any]       =				is_training
									__A	:						Dict       =				use_labels
									__A	:						List[str]       =				hidden_size
									__A	:						Union[str, Any]       =				num_hidden_layers
									__A	:						List[str]       =				num_attention_heads
									__A	:						Tuple       =				intermediate_size
									__A	:						List[str]       =				hidden_act
									__A	:						Dict       =				hidden_dropout_prob
									__A	:						str       =				attention_probs_dropout_prob
									__A	:						List[str]       =				type_sequence_label_size
									__A	:						Union[str, Any]       =				initializer_range
									__A	:						int       =				mask_ratio
									__A	:						List[str]       =				scope

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



				def 					UpperCAmelCase_		(			self     ):
									__A	:						Union[str, Any]       =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

									__A	:						int       =				None
									if self.use_labels:
														__A	:						List[Any]       =				ids_tensor([self.batch_size]     ,					self.type_sequence_label_size     )

									__A	:						Optional[int]       =				self.get_config()

									return config, pixel_values, labels



				def 					UpperCAmelCase_		(			self     ):
									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=_A     ,					initializer_range=self.initializer_range     ,					mask_ratio=self.mask_ratio     ,					)



				def 					UpperCAmelCase_		(			self     ,					_A     ,					_A     ,					_A     ):
									__A	:						Tuple       =				ViTMAEModel(config=_A     )
									model.to(_A     )
									model.eval()
									__A	:						Tuple       =				model(_A     )
									self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, self.seq_length, self.hidden_size)     )



				def 					UpperCAmelCase_		(			self     ,					_A     ,					_A     ,					_A     ):
									__A	:						List[Any]       =				ViTMAEForPreTraining(_A     )
									model.to(_A     )
									model.eval()
									__A	:						Optional[int]       =				model(_A     )
									__A	:						int       =				(self.image_size // self.patch_size) ** 2
									__A	:						List[str]       =				self.patch_size**2 * self.num_channels
									self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, num_patches, expected_num_channels)     )

									# test greyscale images
									__A	:						Tuple       =				1
									__A	:						Any       =				ViTMAEForPreTraining(_A     )
									model.to(_A     )
									model.eval()
									__A	:						Optional[Any]       =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
									__A	:						Dict       =				model(_A     )
									__A	:						List[str]       =				self.patch_size**2
									self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, num_patches, expected_num_channels)     )





				def 					UpperCAmelCase_		(			self     ):
									__A	:						Dict       =				self.prepare_config_and_inputs()
									__A	:						Tuple       =				config_and_inputs
									__A	:						Any       =				{'pixel_values': pixel_values}
									return config, inputs_dict





@require_torch
class 						_A(					_a				, _a				, unittest.TestCase   ):


				"""simple docstring"""


				UpperCamelCase    :  Optional[Any]					    =	(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
				UpperCamelCase    :  str					    =	{"""feature-extraction""": ViTMAEModel} if is_torch_available() else {}

				UpperCamelCase    :  Optional[int]					    =	False
				UpperCamelCase    :  List[str]					    =	False
				UpperCamelCase    :  Dict					    =	False
				UpperCamelCase    :  List[str]					    =	False



				def 					UpperCAmelCase_		(			self     ):
									__A	:						List[Any]       =				ViTMAEModelTester(self     )
									__A	:						Union[str, Any]       =				ConfigTester(self     ,					config_class=_A     ,					has_text_modality=_A     ,					hidden_size=37     )



				def 					UpperCAmelCase_		(			self     ):
									self.config_tester.run_common_tests()



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



				def 					UpperCAmelCase_		(			self     ):
									__A	:						Any       =				self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
														__A	:						Union[str, Any]       =				model_class(_A     )
														self.assertIsInstance(model.get_input_embeddings()     ,					(nn.Module)     )
														__A	:						str       =				model.get_output_embeddings()
														self.assertTrue(x is None or isinstance(_A     ,					nn.Linear     )     )



				def 					UpperCAmelCase_		(			self     ):
									__A	:						Dict       =				self.model_tester.prepare_config_and_inputs_for_common()

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

														__A	:						Optional[Any]       =				['pixel_values']
														self.assertListEqual(arg_names[:1]     ,					_A     )



				def 					UpperCAmelCase_		(			self     ):
									__A	:						str       =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_A     )



				def 					UpperCAmelCase_		(			self     ):
									__A	:						int       =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*_A     )



				def 					UpperCAmelCase_		(			self     ,					_A     ,					_A     ,					_A     ):
									np.random.seed(2     )

									__A	:						Optional[Any]       =				int((pt_model.config.image_size // pt_model.config.patch_size) ** 2     )
									__A	:						Optional[int]       =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)     )
									__A	:						List[str]       =				torch.from_numpy(_A     )

									# Add `noise` argument.
									# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
									__A	:						List[str]       =				pt_noise

									super().check_pt_tf_models(_A     ,					_A     ,					_A     )



				def 					UpperCAmelCase_		(			self     ):
									__A	:						Tuple       =				self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
														__A	:						Any       =				model_class(_A     )
														model.to(_A     )
														model.eval()
														# make random mask reproducible
														torch.manual_seed(2     )
														with torch.no_grad():
																			__A	:						List[str]       =				model(**self._prepare_for_class(_A     ,					_A     )     )

														__A	:						str       =				outputs[0].cpu().numpy()
														__A	:						List[str]       =				0

														with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_pretrained(_A     )
																			__A	:						Tuple       =				model_class.from_pretrained(_A     )
																			model.to(_A     )
																			# make random mask reproducible
																			torch.manual_seed(2     )
																			with torch.no_grad():
																								__A	:						Optional[int]       =				model(**self._prepare_for_class(_A     ,					_A     )     )

																			# Make sure we don't have nans
																			__A	:						Tuple       =				after_outputs[0].cpu().numpy()
																			__A	:						List[Any]       =				0
																			__A	:						Any       =				np.amax(np.abs(out_a - out_a     )     )
																			self.assertLessEqual(_A     ,					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 					UpperCAmelCase_		(			self     ):
									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 					UpperCAmelCase_		(			self     ):
									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 					UpperCAmelCase_		(			self     ):
									pass



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



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





				@slow
				def 					UpperCAmelCase_		(			self     ):

									for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__A	:						Tuple       =				ViTMAEModel.from_pretrained(_A     )
														self.assertIsNotNone(_A     )







def  _SCREAMING_SNAKE_CASE				(   )      ->      Tuple:
					__A	:						int       =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
					return image





@require_torch
@require_vision
class 						_A(					unittest.TestCase   ):


				"""simple docstring"""


				@cached_property
				def 					UpperCAmelCase_		(			self     ):
									return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'     ) if is_vision_available() else None





				@slow
				def 					UpperCAmelCase_		(			self     ):
									np.random.seed(2     )

									__A	:						Dict       =				ViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'     ).to(_A     )

									__A	:						int       =				self.default_image_processor
									__A	:						List[Any]       =				prepare_img()
									__A	:						List[Any]       =				image_processor(images=_A     ,					return_tensors='pt'     ).to(_A     )

									# 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)
									__A	:						Optional[Any]       =				ViTMAEConfig()
									__A	:						Dict       =				int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2     )
									__A	:						Dict       =				np.random.uniform(size=(1, num_patches)     )

									# forward pass
									with torch.no_grad():
														__A	:						List[Any]       =				model(**_A     ,					noise=torch.from_numpy(_A     ).to(device=_A     )     )

									# verify the logits
									__A	:						Optional[Any]       =				torch.Size((1, 196, 768)     )
									self.assertEqual(outputs.logits.shape     ,					_A     )

									__A	:						Any       =				torch.tensor(
									    [[-0.0_5_4_8, -1.7_0_2_3, -0.9_3_2_5], [0.3_7_2_1, -0.5_6_7_0, -0.2_2_3_3], [0.8_2_3_5, -1.3_8_7_8, -0.3_5_2_4]]     )

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








from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
 from PIL import Image

 from ..image_utils import load_image

if is_tf_available():
 import tensorflow as tf

 from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
 from ..tf_utils import stable_softmax

if is_torch_available():
 from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

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






@add_end_docstrings(_a   )
class _lowerCamelCase(						_a   ):
       def __init__(					self,      *lowerCamelCase,      **lowerCamelCase)       ->   int:




           """simple docstring"""

           super().__init__(*lowerCamelCase,      **lowerCamelCase)
           requires_backends(self,      'vision')
           self.check_model_type(
               TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
               if self.framework == 'tf'
               else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING)
       def 			UpperCamelCase     (					self,      lowerCamelCase=None)       ->   int:




           """simple docstring"""

           _lowercase       :							Dict      		=		{}
           if top_k is not None:
               _lowercase       :							List[str]      		=		top_k
           return {}, {}, postprocess_params
       def __call__(					self,      lowerCamelCase,      **lowerCamelCase)       ->   Tuple:




           """simple docstring"""

           return super().__call__(lowerCamelCase,      **lowerCamelCase)
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   str:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		load_image(lowerCamelCase)
           _lowercase       :							List[str]      		=		self.image_processor(images=lowerCamelCase,      return_tensors=self.framework)
           return model_inputs
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		self.model(**lowerCamelCase)
           return model_outputs




       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase=5)       ->   Dict:




           """simple docstring"""

           if top_k > self.model.config.num_labels:
               _lowercase       :							List[Any]      		=		self.model.config.num_labels

           if self.framework == "pt":
               _lowercase       :							int      		=		model_outputs.logits.softmax(-1)[0]
               _lowercase		,				_lowercase       :							Union[str, Any]      		=		probs.topk(lowerCamelCase)
           elif self.framework == "tf":
               _lowercase       :							int      		=		stable_softmax(model_outputs.logits,      axis=-1)[0]
               _lowercase       :							List[Any]      		=		tf.math.top_k(lowerCamelCase,      k=lowerCamelCase)
               _lowercase		,				_lowercase       :							Any      		=		topk.values.numpy(), topk.indices.numpy()
           else:
               raise ValueError(F'''Unsupported framework: {self.framework}''')

           _lowercase       :							str      		=		scores.tolist()
           _lowercase       :							str      		=		ids.tolist()
           return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(lowerCamelCase,      lowerCamelCase)]


