



from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

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 TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class 				lowerCAmelCase_		:
       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_=2,			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_=0.6,			SCREAMING_SNAKE_CASE_=None,			)     ->		Any:
          UpperCamelCase    :  Tuple							=	parent
          UpperCamelCase    :  Optional[Any]							=	batch_size
          UpperCamelCase    :  Union[str, Any]							=	image_size
          UpperCamelCase    :  str							=	patch_size
          UpperCamelCase    :  Optional[Any]							=	num_channels
          UpperCamelCase    :  Any							=	is_training
          UpperCamelCase    :  Dict							=	use_labels
          UpperCamelCase    :  int							=	hidden_size
          UpperCamelCase    :  Dict							=	num_hidden_layers
          UpperCamelCase    :  int							=	num_attention_heads
          UpperCamelCase    :  Any							=	intermediate_size
          UpperCamelCase    :  List[Any]							=	hidden_act
          UpperCamelCase    :  str							=	hidden_dropout_prob
          UpperCamelCase    :  int							=	attention_probs_dropout_prob
          UpperCamelCase    :  List[Any]							=	type_sequence_label_size
          UpperCamelCase    :  Optional[int]							=	initializer_range
          UpperCamelCase    :  Any							=	mask_ratio
          UpperCamelCase    :  Optional[int]							=	scope

          # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
          # (we add 1 for the [CLS] token)
          UpperCamelCase    :  int							=	(image_size // patch_size) ** 2
          UpperCamelCase    :  Optional[Any]							=	int(math.ceil((1 - mask_ratio) * (num_patches + 1)       )       )
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          UpperCamelCase    :  Optional[int]							=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

          UpperCamelCase    :  Union[str, Any]							=	None
          if self.use_labels:
             UpperCamelCase    :  Any							=	ids_tensor([self.batch_size],			self.type_sequence_label_size       )

          UpperCamelCase    :  List[Any]							=	self.get_config()

          return config, pixel_values, labels
       def 					snake_case_       (			self       )     ->		str:
          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,			decoder_hidden_size=self.hidden_size,			decoder_num_hidden_layers=self.num_hidden_layers,			decoder_num_attention_heads=self.num_attention_heads,			decoder_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 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Union[str, Any]:
          UpperCamelCase    :  List[Any]							=	TFViTMAEModel(config=_UpperCamelCase       )
          UpperCamelCase    :  List[Any]							=	model(_UpperCamelCase,			training=_UpperCamelCase       )
          self.parent.assertEqual(result.last_hidden_state.shape,			(self.batch_size, self.seq_length, self.hidden_size)       )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Optional[int]:
          UpperCamelCase    :  Optional[int]							=	TFViTMAEForPreTraining(_UpperCamelCase       )
          UpperCamelCase    :  Tuple							=	model(_UpperCamelCase,			training=_UpperCamelCase       )
          # expected sequence length = num_patches
          UpperCamelCase    :  Any							=	(self.image_size // self.patch_size) ** 2
          UpperCamelCase    :  int							=	self.patch_size**2 * self.num_channels
          self.parent.assertEqual(result.logits.shape,			(self.batch_size, num_patches, expected_num_channels)       )

          # test greyscale images
          UpperCamelCase    :  Tuple							=	1
          UpperCamelCase    :  str							=	TFViTMAEForPreTraining(_UpperCamelCase       )

          UpperCamelCase    :  int							=	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
          UpperCamelCase    :  Optional[int]							=	model(_UpperCamelCase,			training=_UpperCamelCase       )
          UpperCamelCase    :  Optional[int]							=	self.patch_size**2
          self.parent.assertEqual(result.logits.shape,			(self.batch_size, num_patches, expected_num_channels)       )






       def 					snake_case_       (			self       )     ->		Dict:
          UpperCamelCase    :  List[Any]							=	self.prepare_config_and_inputs()
          ((UpperCamelCase)	,     (UpperCamelCase)	,     (UpperCamelCase))    :  int							=	config_and_inputs
          UpperCamelCase    :  Union[str, Any]							=	{'pixel_values': pixel_values}
          return config, inputs_dict






@require_tf
class 				lowerCAmelCase_		(  __A       ,    __A       ,    unittest.TestCase							):
       UpperCAmelCase__				:					str        =  (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
       UpperCAmelCase__				:					Tuple        =  {"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

       UpperCAmelCase__				:					Union[str, Any]        =  False
       UpperCAmelCase__				:					Optional[Any]        =  False
       UpperCAmelCase__				:					Union[str, Any]        =  False
       UpperCAmelCase__				:					Union[str, Any]        =  False
       def 					snake_case_       (			self       )     ->		List[Any]:
          UpperCamelCase    :  Dict							=	TFViTMAEModelTester(self       )
          UpperCamelCase    :  str							=	ConfigTester(self,			config_class=_UpperCamelCase,			has_text_modality=_UpperCamelCase,			hidden_size=37       )
       def 					snake_case_       (			self       )     ->		int:
          self.config_tester.run_common_tests()
       @unittest.skip(reason='ViTMAE does not use inputs_embeds'       )
       def 					snake_case_       (			self       )     ->		int:
          pass
       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase	,     UpperCamelCase    :  Union[str, Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase    :  List[str]							=	model_class(_UpperCamelCase       )
             self.assertIsInstance(model.get_input_embeddings(),			(tf.keras.layers.Layer)       )
             UpperCamelCase    :  Optional[Any]							=	model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_UpperCamelCase,			tf.keras.layers.Layer       )       )
       def 					snake_case_       (			self       )     ->		Union[str, Any]:
          UpperCamelCase	,     UpperCamelCase    :  Optional[Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
             UpperCamelCase    :  Any							=	model_class(_UpperCamelCase       )
             UpperCamelCase    :  Optional[int]							=	inspect.signature(model.call       )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase    :  str							=	[*signature.parameters.keys()]

             UpperCamelCase    :  Tuple							=	['pixel_values']
             self.assertListEqual(arg_names[:1],			_UpperCamelCase       )
       def 					snake_case_       (			self       )     ->		int:
          UpperCamelCase    :  Tuple							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_UpperCamelCase       )
       def 					snake_case_       (			self       )     ->		Dict:
          UpperCamelCase    :  Optional[int]							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*_UpperCamelCase       )
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          # make the mask reproducible
          np.random.seed(2       )

          UpperCamelCase	,     UpperCamelCase    :  str							=	self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase    :  Tuple							=	int((config.image_size // config.patch_size) ** 2       )
          UpperCamelCase    :  str							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
             UpperCamelCase    :  str							=	model_class(_UpperCamelCase       )
             UpperCamelCase    :  Optional[Any]							=	self._prepare_for_class(_UpperCamelCase,			_UpperCamelCase       )

             UpperCamelCase    :  List[Any]							=	model(_UpperCamelCase,			noise=_UpperCamelCase       )

             UpperCamelCase    :  Any							=	copy.deepcopy(self._prepare_for_class(_UpperCamelCase,			_UpperCamelCase       )       )
             UpperCamelCase    :  int							=	model(**_UpperCamelCase,			noise=_UpperCamelCase       )
             UpperCamelCase    :  Union[str, Any]							=	outputs_dict[0].numpy()
             UpperCamelCase    :  Optional[int]							=	outputs_keywords[0].numpy()

             self.assertLess(np.sum(np.abs(output_dict - output_keywords       )       ),			1e-6       )
       def 					snake_case_       (			self       )     ->		Optional[Any]:
          # make the mask reproducible
          np.random.seed(2       )

          UpperCamelCase	,     UpperCamelCase    :  List[str]							=	self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase    :  int							=	int((config.image_size // config.patch_size) ** 2       )
          UpperCamelCase    :  Union[str, Any]							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          def prepare_numpy_arrays(SCREAMING_SNAKE_CASE_       ):
             UpperCamelCase    :  str							=	{}
             for k, v in inputs_dict.items():
                if tf.is_tensor(_UpperCamelCase       ):
                   UpperCamelCase    :  Tuple							=	v.numpy()
                else:
                   UpperCamelCase    :  List[str]							=	np.array(_UpperCamelCase       )

             return inputs_np_dict

          for model_class in self.all_model_classes:
             UpperCamelCase    :  Optional[int]							=	model_class(_UpperCamelCase       )

             UpperCamelCase    :  List[str]							=	self._prepare_for_class(_UpperCamelCase,			_UpperCamelCase       )
             UpperCamelCase    :  List[Any]							=	prepare_numpy_arrays(_UpperCamelCase       )

             UpperCamelCase    :  Any							=	model(_UpperCamelCase,			noise=_UpperCamelCase       )
             UpperCamelCase    :  Dict							=	model(**_UpperCamelCase,			noise=_UpperCamelCase       )
             self.assert_outputs_same(_UpperCamelCase,			_UpperCamelCase       )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       )     ->		Dict:
          # make masks reproducible
          np.random.seed(2       )

          UpperCamelCase    :  Optional[int]							=	int((tf_model.config.image_size // tf_model.config.patch_size) ** 2       )
          UpperCamelCase    :  Tuple							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
          UpperCamelCase    :  Union[str, Any]							=	tf.constant(_UpperCamelCase       )

          # Add `noise` argument.
          # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
          UpperCamelCase    :  Optional[int]							=	tf_noise

          super().check_pt_tf_models(_UpperCamelCase,			_UpperCamelCase,			_UpperCamelCase       )
       def 					snake_case_       (			self       )     ->		str:
          # make mask reproducible
          np.random.seed(2       )

          UpperCamelCase	,     UpperCamelCase    :  str							=	self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase    :  Tuple							=	{
              module_member
              for model_class in self.all_model_classes
              for module in (import_module(model_class.__module__       ),)
              for module_member_name in dir(_UpperCamelCase       )
              if module_member_name.endswith('MainLayer'       )
              # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
              and module_member_name[: -len('MainLayer'       )] == model_class.__name__[: -len('Model'       )]
              for module_member in (getattr(_UpperCamelCase,			_UpperCamelCase       ),)
              if isinstance(_UpperCamelCase,			_UpperCamelCase       )
              and tf.keras.layers.Layer in module_member.__bases__
              and getattr(_UpperCamelCase,			'_keras_serializable',			_UpperCamelCase       )
          }

          UpperCamelCase    :  Tuple							=	int((config.image_size // config.patch_size) ** 2       )
          UpperCamelCase    :  Dict							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )
          UpperCamelCase    :  Dict							=	tf.convert_to_tensor(_UpperCamelCase       )
          inputs_dict.update({'noise': noise}       )

          for main_layer_class in tf_main_layer_classes:
             UpperCamelCase    :  int							=	main_layer_class(_UpperCamelCase       )

             UpperCamelCase    :  str							=	{
                 name: tf.keras.Input(tensor.shape[1:],			dtype=tensor.dtype       ) for name, tensor in inputs_dict.items()
             }

             UpperCamelCase    :  Any							=	tf.keras.Model(_UpperCamelCase,			outputs=main_layer(_UpperCamelCase       )       )
             UpperCamelCase    :  int							=	model(_UpperCamelCase       )

             with tempfile.TemporaryDirectory() as tmpdirname:
                UpperCamelCase    :  Optional[int]							=	os.path.join(_UpperCamelCase,			'keras_model.h5'       )
                model.save(_UpperCamelCase       )
                UpperCamelCase    :  Tuple							=	tf.keras.models.load_model(
                    _UpperCamelCase,			custom_objects={main_layer_class.__name__: main_layer_class}       )
                assert isinstance(_UpperCamelCase,			tf.keras.Model       )
                UpperCamelCase    :  Optional[Any]							=	model(_UpperCamelCase       )
                self.assert_outputs_same(_UpperCamelCase,			_UpperCamelCase       )
       @slow
       def 					snake_case_       (			self       )     ->		str:
          # make mask reproducible
          np.random.seed(2       )

          UpperCamelCase	,     UpperCamelCase    :  Optional[Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase    :  Dict							=	int((config.image_size // config.patch_size) ** 2       )
          UpperCamelCase    :  List[Any]							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
             UpperCamelCase    :  List[str]							=	model_class(_UpperCamelCase       )
             UpperCamelCase    :  Optional[int]							=	self._prepare_for_class(_UpperCamelCase,			_UpperCamelCase       )
             UpperCamelCase    :  List[str]							=	model(_UpperCamelCase,			noise=_UpperCamelCase       )

             if model_class.__name__ == "TFViTMAEModel":
                UpperCamelCase    :  Optional[Any]							=	outputs.last_hidden_state.numpy()
                UpperCamelCase    :  int							=	0
             else:
                UpperCamelCase    :  List[Any]							=	outputs.logits.numpy()
                UpperCamelCase    :  Optional[Any]							=	0

             with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(_UpperCamelCase,			saved_model=_UpperCamelCase       )
                UpperCamelCase    :  List[str]							=	model_class.from_pretrained(_UpperCamelCase       )
                UpperCamelCase    :  Optional[Any]							=	model(_UpperCamelCase,			noise=_UpperCamelCase       )

                if model_class.__name__ == "TFViTMAEModel":
                   UpperCamelCase    :  int							=	after_outputs['last_hidden_state'].numpy()
                   UpperCamelCase    :  Tuple							=	0
                else:
                   UpperCamelCase    :  int							=	after_outputs['logits'].numpy()
                   UpperCamelCase    :  Any							=	0

                UpperCamelCase    :  Union[str, Any]							=	np.amax(np.abs(out_a - out_a       )       )
                self.assertLessEqual(_UpperCamelCase,			1e-5       )
       def 					snake_case_       (			self       )     ->		List[str]:
          # make mask reproducible
          np.random.seed(2       )

          UpperCamelCase	,     UpperCamelCase    :  Optional[Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

          UpperCamelCase    :  Optional[int]							=	int((config.image_size // config.patch_size) ** 2       )
          UpperCamelCase    :  Tuple							=	np.random.uniform(size=(self.model_tester.batch_size, num_patches)       )

          for model_class in self.all_model_classes:
             UpperCamelCase    :  List[Any]							=	model_class(_UpperCamelCase       )
             UpperCamelCase    :  str							=	self._prepare_for_class(_UpperCamelCase,			_UpperCamelCase       )

             UpperCamelCase    :  Union[str, Any]							=	model(_UpperCamelCase,			noise=_UpperCamelCase       )
             UpperCamelCase    :  Optional[Any]							=	model.get_config()
             # make sure that returned config is jsonifiable, which is required by keras
             json.dumps(_UpperCamelCase       )
             UpperCamelCase    :  Union[str, Any]							=	model_class.from_config(model.get_config()       )
             # make sure it also accepts a normal config
             UpperCamelCase    :  Optional[int]							=	model_class.from_config(model.config       )
             UpperCamelCase    :  Optional[Any]							=	new_model(_UpperCamelCase       )  # Build model
             new_model.set_weights(model.get_weights()       )
             UpperCamelCase    :  List[str]							=	new_model(_UpperCamelCase,			noise=_UpperCamelCase       )

             self.assert_outputs_same(_UpperCamelCase,			_UpperCamelCase       )
       @unittest.skip(
           reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'       )
       def 					snake_case_       (			self       )     ->		List[Any]:
          pass
       @unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'       )
       def 					snake_case_       (			self       )     ->		Optional[int]:
          pass






       @slow
       def 					snake_case_       (			self       )     ->		List[str]:
          UpperCamelCase    :  Any							=	TFViTMAEModel.from_pretrained('google/vit-base-patch16-224'       )
          self.assertIsNotNone(_UpperCamelCase       )






def 		UpperCamelCase	(   )	->     List[str]:
   UpperCamelCase    :  Optional[int]							=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
   return image



@require_tf
@require_vision
class 				lowerCAmelCase_		(  unittest.TestCase							):
       @cached_property
       def 					snake_case_       (			self       )     ->		Any:
          return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'       ) if is_vision_available() else None






       @slow
       def 					snake_case_       (			self       )     ->		str:
          # make random mask reproducible across the PT and TF model
          np.random.seed(2       )

          UpperCamelCase    :  List[Any]							=	TFViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'       )

          UpperCamelCase    :  int							=	self.default_image_processor
          UpperCamelCase    :  int							=	prepare_img()
          UpperCamelCase    :  List[Any]							=	image_processor(images=_UpperCamelCase,			return_tensors='tf'       )

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

          # forward pass
          UpperCamelCase    :  Union[str, Any]							=	model(**_UpperCamelCase,			noise=_UpperCamelCase       )

          # verify the logits
          UpperCamelCase    :  Tuple							=	tf.convert_to_tensor([1, 196, 768]       )
          self.assertEqual(outputs.logits.shape,			_UpperCamelCase       )

          UpperCamelCase    :  List[Any]							=	tf.convert_to_tensor(
              [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]       )

          tf.debugging.assert_near(outputs.logits[0, :3, :3],			_UpperCamelCase,			atol=1e-4       )






from collections import deque

from .hash_table import HashTable





class 	snake_case_   (		__A					):







			'''simple docstring'''

			def __init__(     self :				int							,   *_UpperCamelCase :				int							,   **_UpperCamelCase :				Tuple							)      ->Tuple:
									super().__init__(*_UpperCamelCase							,   **_UpperCamelCase							)







			def 				snake_case__(     self :				Dict							,   _UpperCamelCase :				List[str]							,   _UpperCamelCase :				Dict							)      ->Tuple:
									snake_case_   =	deque([]							) if self.values[key] is None else self.values[key]
									self.values[key].appendleft(_UpperCamelCase							)
									snake_case_   =	self.values[key]







			def 				snake_case__(     self :				List[Any]							)      ->str:
									return (
									    sum(self.charge_factor - len(_UpperCamelCase							) for slot in self.values							)
									    / self.size_table
									    * self.charge_factor
									)







			def 				snake_case__(     self :				Dict							,   _UpperCamelCase :				Union[str, Any]							,   _UpperCamelCase :				Optional[int]=None							)      ->str:
									if not (
									    len(self.values[key]							) == self.charge_factor and self.values.count(_UpperCamelCase							) == 0
									):
															return key
									return super()._collision_resolution(_UpperCamelCase							,   _UpperCamelCase							)