'''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						:						Dict    ,							a						:						List[Any]    ,							a						:						Optional[Any]=13    ,							a						:						Dict=30    ,							a						:						Union[str, Any]=2    ,							a						:						Dict=3    ,							a						:						Optional[int]=True    ,							a						:						Optional[Any]=True    ,							a						:						Union[str, Any]=32    ,							a						:						int=2    ,							a						:						List[Any]=4    ,							a						:						List[Any]=37    ,							a						:						Tuple="gelu"    ,							a						:						Optional[Any]=0.1    ,							a						:						List[Any]=0.1    ,							a						:						Tuple=10    ,							a						:						List[str]=0.02    ,							a						:						str=3    ,							a						:						str=0.6    ,							a						:						Any=None    ,							):





         """simple docstring"""



         __lowerCamelCase            =    parent
         __lowerCamelCase            =    batch_size
         __lowerCamelCase            =    image_size
         __lowerCamelCase            =    patch_size
         __lowerCamelCase            =    num_channels
         __lowerCamelCase            =    is_training
         __lowerCamelCase            =    use_labels
         __lowerCamelCase            =    hidden_size
         __lowerCamelCase            =    num_hidden_layers
         __lowerCamelCase            =    num_attention_heads
         __lowerCamelCase            =    intermediate_size
         __lowerCamelCase            =    hidden_act
         __lowerCamelCase            =    hidden_dropout_prob
         __lowerCamelCase            =    attention_probs_dropout_prob
         __lowerCamelCase            =    type_sequence_label_size
         __lowerCamelCase            =    initializer_range
         __lowerCamelCase            =    mask_ratio
         __lowerCamelCase            =    scope

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






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

         __lowerCamelCase            =    self.get_config()

         return config, pixel_values, labels






     def   SCREAMING_SNAKE_CASE__  (					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=a    ,							initializer_range=self.initializer_range    ,							mask_ratio=self.mask_ratio    ,							)






     def   SCREAMING_SNAKE_CASE__  (					self						:						Any    ,							a						:						Optional[int]    ,							a						:						Union[str, Any]    ,							a						:						Union[str, Any]       ):





         """simple docstring"""



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






     def   SCREAMING_SNAKE_CASE__  (					self						:						Dict    ,							a						:						Tuple    ,							a						:						Dict    ,							a						:						List[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    TFViTMAEForPreTraining(a       )
         __lowerCamelCase            =    model(a    ,							training=a       )
         # expected sequence length = num_patches
         __lowerCamelCase            =    (self.image_size // self.patch_size) ** 2
         __lowerCamelCase            =    self.patch_size**2 * self.num_channels
         self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, num_patches, expected_num_channels)       )

         # test greyscale images
         __lowerCamelCase            =    1
         __lowerCamelCase            =    TFViTMAEForPreTraining(a       )

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






     def   SCREAMING_SNAKE_CASE__  (					self						:						Union[str, Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    self.prepare_config_and_inputs()
         ((__lowerCamelCase)      ,    (__lowerCamelCase)      ,    (__lowerCamelCase))            =    config_and_inputs
         __lowerCamelCase            =    {'''pixel_values''': pixel_values}
         return config, inputs_dict




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

     lowerCamelCase    :	Union[str, Any]				=False
     lowerCamelCase    :	Optional[Any]				=False
     lowerCamelCase    :	Optional[int]				=False
     lowerCamelCase    :	Union[str, Any]				=False






     def   SCREAMING_SNAKE_CASE__  (					self						:						int       ):





         """simple docstring"""



         __lowerCamelCase            =    TFViTMAEModelTester(self       )
         __lowerCamelCase            =    ConfigTester(self    ,							config_class=a    ,							has_text_modality=a    ,							hidden_size=37       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						str       ):





         """simple docstring"""



         self.config_tester.run_common_tests()






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





         """simple docstring"""



         pass






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[str]       ):





         """simple docstring"""



         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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






     def   SCREAMING_SNAKE_CASE__  (					self						:						Dict       ):





         """simple docstring"""



         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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

             __lowerCamelCase            =    ['''pixel_values''']
             self.assertListEqual(arg_names[:1]    ,							a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[str]       ):





         """simple docstring"""



         __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						Tuple       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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

         for model_class in self.all_model_classes:
             __lowerCamelCase            =    model_class(a       )
             __lowerCamelCase            =    self._prepare_for_class(a    ,							a       )

             __lowerCamelCase            =    model(a    ,							noise=a       )

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

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






     def   SCREAMING_SNAKE_CASE__  (					self						:						List[str]       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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

         def prepare_numpy_arrays(a						:						Optional[Any]       ):
             __lowerCamelCase            =    {}
             for k, v in inputs_dict.items():
                 if tf.is_tensor(a       ):
                     __lowerCamelCase            =    v.numpy()
                 else:
                     __lowerCamelCase            =    np.array(a       )

             return inputs_np_dict

         for model_class in self.all_model_classes:
             __lowerCamelCase            =    model_class(a       )

             __lowerCamelCase            =    self._prepare_for_class(a    ,							a       )
             __lowerCamelCase            =    prepare_numpy_arrays(a       )

             __lowerCamelCase            =    model(a    ,							noise=a       )
             __lowerCamelCase            =    model(**a    ,							noise=a       )
             self.assert_outputs_same(a    ,							a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						int    ,							a						:						Optional[int]    ,							a						:						List[str]    ,							a						:						Tuple       ):





         """simple docstring"""



         np.random.seed(2       )

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

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

         super().check_pt_tf_models(a    ,							a    ,							a       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[int]       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

         __lowerCamelCase            =    {
             module_member
             for model_class in self.all_model_classes
             for module in (import_module(model_class.__module__       ),)
             for module_member_name in dir(a       )
             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(a    ,							a       ),)
             if isinstance(a    ,							a       )
             and tf.keras.layers.Layer in module_member.__bases__
             and getattr(a    ,							'''_keras_serializable'''    ,							a       )
         }

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

         for main_layer_class in tf_main_layer_classes:
             __lowerCamelCase            =    main_layer_class(a       )

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

             __lowerCamelCase            =    tf.keras.Model(a    ,							outputs=main_layer(a       )       )
             __lowerCamelCase            =    model(a       )

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






     @slow
     def   SCREAMING_SNAKE_CASE__  (					self						:						str       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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

         for model_class in self.all_model_classes:
             __lowerCamelCase            =    model_class(a       )
             __lowerCamelCase            =    self._prepare_for_class(a    ,							a       )
             __lowerCamelCase            =    model(a    ,							noise=a       )

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

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

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

                 __lowerCamelCase            =    np.amax(np.abs(out_a - out_a       )       )
                 self.assertLessEqual(a    ,							1e-5       )






     def   SCREAMING_SNAKE_CASE__  (					self						:						int       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase      ,    __lowerCamelCase            =    self.model_tester.prepare_config_and_inputs_for_common()

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

         for model_class in self.all_model_classes:
             __lowerCamelCase            =    model_class(a       )
             __lowerCamelCase            =    self._prepare_for_class(a    ,							a       )

             __lowerCamelCase            =    model(a    ,							noise=a       )
             __lowerCamelCase            =    model.get_config()
             # make sure that returned config is jsonifiable, which is required by keras
             json.dumps(a       )
             __lowerCamelCase            =    model_class.from_config(model.get_config()       )
             # make sure it also accepts a normal config
             __lowerCamelCase            =    model_class.from_config(model.config       )
             __lowerCamelCase            =    new_model(a       )  # Build model
             new_model.set_weights(model.get_weights()       )
             __lowerCamelCase            =    new_model(a    ,							noise=a       )

             self.assert_outputs_same(a    ,							a       )






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





         """simple docstring"""



         pass






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





         """simple docstring"""



         pass






     @slow
     def   SCREAMING_SNAKE_CASE__  (					self						:						str       ):





         """simple docstring"""



         __lowerCamelCase            =    TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224'''       )
         self.assertIsNotNone(a       )






def __lowerCAmelCase					(				)      ->       Dict:
    __lowerCamelCase            =    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
    return image




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






     @cached_property
     def   SCREAMING_SNAKE_CASE__  (					self						:						Union[str, Any]       ):





         """simple docstring"""



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






     @slow
     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[int]       ):





         """simple docstring"""



         np.random.seed(2       )

         __lowerCamelCase            =    TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base'''       )

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

         # forward pass
         __lowerCamelCase            =    model(**a    ,							noise=a       )

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

         __lowerCamelCase            =    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]    ,							a    ,							atol=1e-4       )

'''simple docstring'''




def __lowerCAmelCase					(				UpperCamelCase__							)      ->       float:
    __lowerCamelCase            =    0
    while len(UpperCamelCase__							) > 1:
        __lowerCamelCase            =    0
        # Consider two files with minimum cost to be merged
        for _ in range(2							):
            __lowerCamelCase            =    files.index(min(UpperCamelCase__							)							)
            temp += files[min_index]
            files.pop(UpperCamelCase__							)
        files.append(UpperCamelCase__							)
        optimal_merge_cost += temp
    return optimal_merge_cost


if __name__ == "__main__":
  import doctest

  doctest.testmod()
