

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     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase=13      , _UpperCamelCase=30      , _UpperCamelCase=2      , _UpperCamelCase=3      , _UpperCamelCase=True      , _UpperCamelCase=True      , _UpperCamelCase=32      , _UpperCamelCase=2      , _UpperCamelCase=4      , _UpperCamelCase=37      , _UpperCamelCase="gelu"      , _UpperCamelCase=0.1      , _UpperCamelCase=0.1      , _UpperCamelCase=10      , _UpperCamelCase=0.02      , _UpperCamelCase=3      , _UpperCamelCase=0.6      , _UpperCamelCase=None      , )      ->							Optional[int]:






              """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 			a		(		self   )      ->							str:






              """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 			a		(		self   )      ->							Union[str, Any]:






              """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      , 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 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Dict:






              """simple docstring"""


              __snake_case       	=							TFViTMAEModel(config=_UpperCamelCase   )
              __snake_case       	=							model(_UpperCamelCase      , training=_UpperCamelCase   )
              self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Dict:






              """simple docstring"""


              __snake_case       	=							TFViTMAEForPreTraining(_UpperCamelCase   )
              __snake_case       	=							model(_UpperCamelCase      , training=_UpperCamelCase   )
              # expected sequence length = num_patches
              __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       	=							TFViTMAEForPreTraining(_UpperCamelCase   )

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






       def 			a		(		self   )      ->							Tuple:






              """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     __snake_case    (				snake_case__			,	snake_case__			,	unittest.TestCase  ):
       """simple docstring"""







       __SCREAMING_SNAKE_CASE				    =							(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
       __SCREAMING_SNAKE_CASE				    =							{"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

       __SCREAMING_SNAKE_CASE				    =							False
       __SCREAMING_SNAKE_CASE				    =							False
       __SCREAMING_SNAKE_CASE				    =							False
       __SCREAMING_SNAKE_CASE				    =							False


       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


              __snake_case       	=							TFViTMAEModelTester(self   )
              __snake_case       	=							ConfigTester(self      , config_class=_UpperCamelCase      , has_text_modality=_UpperCamelCase      , hidden_size=37   )


       def 			a		(		self   )      ->							Optional[Any]:






              """simple docstring"""


              self.config_tester.run_common_tests()


       @unittest.skip(reason="""ViTMAE does not use inputs_embeds"""   )
       def 			a		(		self   )      ->							List[Any]:






              """simple docstring"""


              pass


       def 			a		(		self   )      ->							List[Any]:






              """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()      , (tf.keras.layers.Layer)   )
                     __snake_case       	=							model.get_output_embeddings()
                     self.assertTrue(x is None or isinstance(_UpperCamelCase      , tf.keras.layers.Layer   )   )


       def 			a		(		self   )      ->							Dict:






              """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.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]      , _UpperCamelCase   )


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*_UpperCamelCase   )


       def 			a		(		self   )      ->							int:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_pretraining(*_UpperCamelCase   )


       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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

              for model_class in self.all_model_classes:
                     __snake_case       	=							model_class(_UpperCamelCase   )
                     __snake_case       	=							self._prepare_for_class(_UpperCamelCase      , _UpperCamelCase   )

                     __snake_case       	=							model(_UpperCamelCase      , noise=_UpperCamelCase   )

                     __snake_case       	=							copy.deepcopy(self._prepare_for_class(_UpperCamelCase      , _UpperCamelCase   )   )
                     __snake_case       	=							model(**_UpperCamelCase      , noise=_UpperCamelCase   )
                     __snake_case       	=							outputs_dict[0].numpy()
                     __snake_case       	=							outputs_keywords[0].numpy()

                     self.assertLess(np.sum(np.abs(output_dict - output_keywords   )   )      , 1E-6   )


       def 			a		(		self   )      ->							List[str]:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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

              def prepare_numpy_arrays(_UpperCamelCase   ):
                     __snake_case       	=							{}
                     for k, v in inputs_dict.items():
                            if tf.is_tensor(_UpperCamelCase   ):
                                   __snake_case       	=							v.numpy()
                            else:
                                   __snake_case       	=							np.array(_UpperCamelCase   )

                     return inputs_np_dict

              for model_class in self.all_model_classes:
                     __snake_case       	=							model_class(_UpperCamelCase   )

                     __snake_case       	=							self._prepare_for_class(_UpperCamelCase      , _UpperCamelCase   )
                     __snake_case       	=							prepare_numpy_arrays(_UpperCamelCase   )

                     __snake_case       	=							model(_UpperCamelCase      , noise=_UpperCamelCase   )
                     __snake_case       	=							model(**_UpperCamelCase      , noise=_UpperCamelCase   )
                     self.assert_outputs_same(_UpperCamelCase      , _UpperCamelCase   )


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Dict:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case       	=							int((tf_model.config.image_size // tf_model.config.patch_size) ** 2   )
              __snake_case       	=							np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )
              __snake_case       	=							tf.constant(_UpperCamelCase   )

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

              super().check_pt_tf_models(_UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )


       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

              __snake_case       	=							{
                  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   )
              }

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

              for main_layer_class in tf_main_layer_classes:
                     __snake_case       	=							main_layer_class(_UpperCamelCase   )

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

                     __snake_case       	=							tf.keras.Model(_UpperCamelCase      , outputs=main_layer(_UpperCamelCase   )   )
                     __snake_case       	=							model(_UpperCamelCase   )

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            __snake_case       	=							os.path.join(_UpperCamelCase      , """keras_model.h5"""   )
                            model.save(_UpperCamelCase   )
                            __snake_case       	=							tf.keras.models.load_model(
                                _UpperCamelCase      , custom_objects={main_layer_class.__name__: main_layer_class}   )
                            assert isinstance(_UpperCamelCase      , tf.keras.Model   )
                            __snake_case       	=							model(_UpperCamelCase   )
                            self.assert_outputs_same(_UpperCamelCase      , _UpperCamelCase   )


       @slow
       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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

              for model_class in self.all_model_classes:
                     __snake_case       	=							model_class(_UpperCamelCase   )
                     __snake_case       	=							self._prepare_for_class(_UpperCamelCase      , _UpperCamelCase   )
                     __snake_case       	=							model(_UpperCamelCase      , noise=_UpperCamelCase   )

                     if model_class.__name__ == "TFViTMAEModel":
                            __snake_case       	=							outputs.last_hidden_state.numpy()
                            __snake_case       	=							0
                     else:
                            __snake_case       	=							outputs.logits.numpy()
                            __snake_case       	=							0

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            model.save_pretrained(_UpperCamelCase      , saved_model=_UpperCamelCase   )
                            __snake_case       	=							model_class.from_pretrained(_UpperCamelCase   )
                            __snake_case       	=							model(_UpperCamelCase      , noise=_UpperCamelCase   )

                            if model_class.__name__ == "TFViTMAEModel":
                                   __snake_case       	=							after_outputs["""last_hidden_state"""].numpy()
                                   __snake_case       	=							0
                            else:
                                   __snake_case       	=							after_outputs["""logits"""].numpy()
                                   __snake_case       	=							0

                            __snake_case       	=							np.amax(np.abs(out_a - out_a   )   )
                            self.assertLessEqual(_UpperCamelCase      , 1E-5   )


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case	,    __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_common()

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

              for model_class in self.all_model_classes:
                     __snake_case       	=							model_class(_UpperCamelCase   )
                     __snake_case       	=							self._prepare_for_class(_UpperCamelCase      , _UpperCamelCase   )

                     __snake_case       	=							model(_UpperCamelCase      , noise=_UpperCamelCase   )
                     __snake_case       	=							model.get_config()
                     # make sure that returned config is jsonifiable, which is required by keras
                     json.dumps(_UpperCamelCase   )
                     __snake_case       	=							model_class.from_config(model.get_config()   )
                     # make sure it also accepts a normal config
                     __snake_case       	=							model_class.from_config(model.config   )
                     __snake_case       	=							new_model(_UpperCamelCase   )  # Build model
                     new_model.set_weights(model.get_weights()   )
                     __snake_case       	=							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
    to get deterministic results."""   )
       def 			a		(		self   )      ->							List[Any]:






              """simple docstring"""


              pass


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






              """simple docstring"""


              pass






       @slow
       def 			a		(		self   )      ->							Tuple:






              """simple docstring"""


              __snake_case       	=							TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224"""   )
              self.assertIsNotNone(_UpperCamelCase   )







def 						lowerCamelCase__      ( ):


       """simple docstring"""




       __snake_case       	=							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
       return image




@require_tf
@require_vision
class     __snake_case    (				unittest.TestCase  ):
       """simple docstring"""







       @cached_property
       def 			a		(		self   )      ->							List[str]:






              """simple docstring"""


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






       @slow
       def 			a		(		self   )      ->							List[str]:






              """simple docstring"""


              np.random.seed(2   )

              __snake_case       	=							TFViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""   )

              __snake_case       	=							self.default_image_processor
              __snake_case       	=							prepare_img()
              __snake_case       	=							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)
              __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
              __snake_case       	=							model(**_UpperCamelCase      , noise=_UpperCamelCase   )

              # verify the logits
              __snake_case       	=							tf.convert_to_tensor([1, 1_96, 7_68]   )
              self.assertEqual(outputs.logits.shape      , _UpperCamelCase   )

              __snake_case       	=							tf.convert_to_tensor(
                  [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]   )

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




def 						lowerCamelCase__      ( __A    :int	,__A    :int		):


       """simple docstring"""




       if a < 0 or b < 0:
              raise ValueError("""the value of both inputs must be positive"""		)

       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"
       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"

       __snake_case       	=							max(len(__A		)	,len(__A		)		)

       return "0b" + "".join(
           str(int(char_a == """1""" and char_b == """1"""		)		)
           for char_a, char_b in zip(a_binary.zfill(__A		)	,b_binary.zfill(__A		)		)		)


if __name__ == "__main__":
    import doctest

    doctest.testmod()

