

'''simple docstring'''







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



    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						,			)  -> str:





           '''simple docstring'''
           A_		   = parent
           A_		   = batch_size
           A_		   = image_size
           A_		   = patch_size
           A_		   = num_channels
           A_		   = is_training
           A_		   = use_labels
           A_		   = hidden_size
           A_		   = num_hidden_layers
           A_		   = num_attention_heads
           A_		   = intermediate_size
           A_		   = hidden_act
           A_		   = hidden_dropout_prob
           A_		   = attention_probs_dropout_prob
           A_		   = type_sequence_label_size
           A_		   = initializer_range
           A_		   = mask_ratio
           A_		   = scope

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



    def    snake_case_     (   self    )  -> Optional[int]:





           '''simple docstring'''
           A_		   = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

           A_		   = self.get_config()

           return config, pixel_values, labels



    def    snake_case_     (   self    )  -> Optional[int]:





           '''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    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )  -> str:





           '''simple docstring'''
           A_		   = TFViTMAEModel(config=UpperCamelCase__    )
           A_		   = 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						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )  -> Union[str, Any]:





           '''simple docstring'''
           A_		   = TFViTMAEForPreTraining(UpperCamelCase__    )
           A_		   = model(UpperCamelCase__						,			training=UpperCamelCase__    )
           # expected sequence length = num_patches
           A_		   = (self.image_size // self.patch_size) ** 2
           A_		   = 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_		   = 1
           A_		   = TFViTMAEForPreTraining(UpperCamelCase__    )

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







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





           '''simple docstring'''
           A_		   = self.prepare_config_and_inputs()
           ((A_)   ,     (A_)   ,     (A_))		   = config_and_inputs
           A_		   = {"""pixel_values""": pixel_values}
           return config, inputs_dict


@require_tf
class    A__ ( _snake_case  ,  _snake_case  ,  unittest.TestCase     ):
    lowercase		   =						(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
    lowercase		   =						{"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

    lowercase		   =						False
    lowercase		   =						False
    lowercase		   =						False
    lowercase		   =						False



    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           A_		   = TFViTMAEModelTester(self    )
           A_		   = ConfigTester(self						,			config_class=UpperCamelCase__						,			has_text_modality=UpperCamelCase__						,			hidden_size=37    )



    def    snake_case_     (   self    )  -> Union[str, Any]:





           '''simple docstring'''
           self.config_tester.run_common_tests()



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





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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



    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

                  A_		   = ["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]						,			UpperCamelCase__    )



    def    snake_case_     (   self    )  -> str:





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__    )



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





           '''simple docstring'''
           A_		   = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*UpperCamelCase__    )



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





           '''simple docstring'''
           # make the mask reproducible
           np.random.seed(2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

           for model_class in self.all_model_classes:
                  A_		   = model_class(UpperCamelCase__    )
                  A_		   = self._prepare_for_class(UpperCamelCase__						,			UpperCamelCase__    )

                  A_		   = model(UpperCamelCase__						,			noise=UpperCamelCase__    )

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

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



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





           '''simple docstring'''
           # make the mask reproducible
           np.random.seed(2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

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

                  return inputs_np_dict

           for model_class in self.all_model_classes:
                  A_		   = model_class(UpperCamelCase__    )

                  A_		   = self._prepare_for_class(UpperCamelCase__						,			UpperCamelCase__    )
                  A_		   = prepare_numpy_arrays(UpperCamelCase__    )

                  A_		   = model(UpperCamelCase__						,			noise=UpperCamelCase__    )
                  A_		   = model(**UpperCamelCase__						,			noise=UpperCamelCase__    )
                  self.assert_outputs_same(UpperCamelCase__						,			UpperCamelCase__    )



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )  -> Union[str, Any]:





           '''simple docstring'''
           # make masks reproducible
           np.random.seed(2    )

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

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

           super().check_pt_tf_models(UpperCamelCase__						,			UpperCamelCase__						,			UpperCamelCase__    )



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





           '''simple docstring'''
           # make mask reproducible
           np.random.seed(2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

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

           for main_layer_class in tf_main_layer_classes:
                  A_		   = main_layer_class(UpperCamelCase__    )

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

                  A_		   = tf.keras.Model(UpperCamelCase__						,			outputs=main_layer(UpperCamelCase__    )    )
                  A_		   = model(UpperCamelCase__    )

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



    @slow
    def    snake_case_     (   self    )  -> Tuple:





           '''simple docstring'''
           # make mask reproducible
           np.random.seed(2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

           for model_class in self.all_model_classes:
                  A_		   = model_class(UpperCamelCase__    )
                  A_		   = self._prepare_for_class(UpperCamelCase__						,			UpperCamelCase__    )
                  A_		   = model(UpperCamelCase__						,			noise=UpperCamelCase__    )

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

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

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

                         A_		   = np.amax(np.abs(out_a - out_a    )    )
                         self.assertLessEqual(UpperCamelCase__						,			1e-5    )



    def    snake_case_     (   self    )  -> str:





           '''simple docstring'''
           # make mask reproducible
           np.random.seed(2    )

           A_   ,     A_		   = self.model_tester.prepare_config_and_inputs_for_common()

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

           for model_class in self.all_model_classes:
                  A_		   = model_class(UpperCamelCase__    )
                  A_		   = self._prepare_for_class(UpperCamelCase__						,			UpperCamelCase__    )

                  A_		   = model(UpperCamelCase__						,			noise=UpperCamelCase__    )
                  A_		   = model.get_config()
                  # make sure that returned config is jsonifiable, which is required by keras
                  json.dumps(UpperCamelCase__    )
                  A_		   = model_class.from_config(model.get_config()    )
                  # make sure it also accepts a normal config
                  A_		   = model_class.from_config(model.config    )
                  A_		   = new_model(UpperCamelCase__    )  # Build model
                  new_model.set_weights(model.get_weights()    )
                  A_		   = 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    snake_case_     (   self    )  -> Optional[Any]:





           '''simple docstring'''
           pass



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





           '''simple docstring'''
           pass







    @slow
    def    snake_case_     (   self    )  -> Tuple:





           '''simple docstring'''
           A_		   = TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224"""    )
           self.assertIsNotNone(UpperCamelCase__    )

def  UpperCAmelCase__      ( )     ->     Optional[int]:
       A_		   = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
       return image







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



    @cached_property
    def    snake_case_     (   self    )  -> str:





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







    @slow
    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           # make random mask reproducible across the PT and TF model
           np.random.seed(2    )

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

           A_		   = self.default_image_processor
           A_		   = prepare_img()
           A_		   = 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)
           A_		   = ViTMAEConfig()
           A_		   = int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2    )
           A_		   = np.random.uniform(size=(1, num_patches)    )

           # forward pass
           A_		   = model(**UpperCamelCase__						,			noise=UpperCamelCase__    )

           # verify the logits
           A_		   = tf.convert_to_tensor([1, 196, 768]    )
           self.assertEqual(outputs.logits.shape						,			UpperCamelCase__    )

           A_		   = 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    )



'''simple docstring'''







import json
import logging
import os
import socket

import git
import numpy as np
import torch


logging.basicConfig(
    format='''%(asctime)s - %(levelname)s - %(name)s - PID: %(process)d -  %(message)s''',
    datefmt='''%m/%d/%Y %H:%M:%S''',
    level=logging.INFO,
)
__lowerCamelCase						    =      logging.getLogger(__name__)
def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     int:
       A_		   = git.Repo(search_parent_directories=UpperCAmelCase__			)
       A_		   = {
           """repo_id""": str(UpperCAmelCase__			),
           """repo_sha""": str(repo.head.object.hexsha			),
           """repo_branch""": str(repo.active_branch			),
       }

       with open(os.path.join(UpperCAmelCase__,    """git_log.json"""			),    """w"""			) as f:
              json.dump(UpperCAmelCase__,    UpperCAmelCase__,    indent=4			)
def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     Dict:
       if params.n_gpu <= 0:
              A_		   = 0
              A_		   = -1
              A_		   = True
              A_		   = False
              return

       assert torch.cuda.is_available()

       logger.info("""Initializing GPUs"""			)
       if params.n_gpu > 1:
              assert params.local_rank != -1

              A_		   = int(os.environ["""WORLD_SIZE"""]			)
              A_		   = int(os.environ["""N_GPU_NODE"""]			)
              A_		   = int(os.environ["""RANK"""]			)

              # number of nodes / node ID
              A_		   = params.world_size // params.n_gpu_per_node
              A_		   = params.global_rank // params.n_gpu_per_node
              A_		   = True

              assert params.n_nodes == int(os.environ["""N_NODES"""]			)
              assert params.node_id == int(os.environ["""NODE_RANK"""]			)

       # local job (single GPU)
       else:
              assert params.local_rank == -1

              A_		   = 1
              A_		   = 0
              A_		   = 0
              A_		   = 0
              A_		   = 1
              A_		   = 1
              A_		   = False

       # sanity checks
       assert params.n_nodes >= 1
       assert 0 <= params.node_id < params.n_nodes
       assert 0 <= params.local_rank <= params.global_rank < params.world_size
       assert params.world_size == params.n_nodes * params.n_gpu_per_node

       # define whether this is the master process / if we are in multi-node distributed mode
       A_		   = params.node_id == 0 and params.local_rank == 0
       A_		   = params.n_nodes > 1

       # summary
       A_		   = F'''--- Global rank: {params.global_rank} - '''
       logger.info(PREFIX + """Number of nodes: %i""" % params.n_nodes			)
       logger.info(PREFIX + """Node ID        : %i""" % params.node_id			)
       logger.info(PREFIX + """Local rank     : %i""" % params.local_rank			)
       logger.info(PREFIX + """World size     : %i""" % params.world_size			)
       logger.info(PREFIX + """GPUs per node  : %i""" % params.n_gpu_per_node			)
       logger.info(PREFIX + """Master         : %s""" % str(params.is_master			)			)
       logger.info(PREFIX + """Multi-node     : %s""" % str(params.multi_node			)			)
       logger.info(PREFIX + """Multi-GPU      : %s""" % str(params.multi_gpu			)			)
       logger.info(PREFIX + """Hostname       : %s""" % socket.gethostname()			)

       # set GPU device
       torch.cuda.set_device(params.local_rank			)

       # initialize multi-GPU
       if params.multi_gpu:
              logger.info("""Initializing PyTorch distributed"""			)
              torch.distributed.init_process_group(
                  init_method="""env://""",    backend="""nccl""",    )







def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     Tuple:
       np.random.seed(args.seed			)
       torch.manual_seed(args.seed			)
       if args.n_gpu > 0:
              torch.cuda.manual_seed_all(args.seed			)
