





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


      def __init__(self	:				str    ,      a__	:				Dict    ,      a__	:				Dict=13    ,      a__	:				Optional[Any]=30    ,      a__	:				int=2    ,      a__	:				Tuple=3    ,      a__	:				str=True    ,      a__	:				Union[str, Any]=True    ,      a__	:				Any=32    ,      a__	:				str=2    ,      a__	:				List[str]=4    ,      a__	:				List[str]=37    ,      a__	:				Optional[Any]="gelu"    ,      a__	:				Any=0.1    ,      a__	:				Dict=0.1    ,      a__	:				Dict=10    ,      a__	:				Optional[Any]=0.0_2    ,      a__	:				List[str]=3    ,      a__	:				str=0.6    ,      a__	:				Union[str, Any]=None    ,      ):





         """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	:				Tuple				):





         """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	:				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        a			(self	:				List[str]    ,      a__	:				Optional[Any]    ,      a__	:				Tuple    ,      a__	:				Dict				):





         """simple docstring"""



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


      def        a			(self	:				List[Any]    ,      a__	:				Any    ,      a__	:				Any    ,      a__	:				Any				):





         """simple docstring"""



         __snake_case								=       TFViTMAEForPreTraining(a__				)
         __snake_case								=       model(a__    ,      training=a__				)
         # 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(a__				)

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


      def        a			(self	:				List[Any]				):





         """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 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    List[str]							=   (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
      A_						:    Tuple							=   {'feature-extraction': TFViTMAEModel} if is_tf_available() else {}

      A_						:    Dict							=   False
      A_						:    Optional[Any]							=   False
      A_						:    List[Any]							=   False
      A_						:    Optional[int]							=   False


      def        a			(self	:				List[Any]				):





         """simple docstring"""



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


      def        a			(self	:				Tuple				):





         """simple docstring"""



         self.config_tester.run_common_tests()


      @unittest.skip(reason='''ViTMAE does not use inputs_embeds'''				)
      def        a			(self	:				List[str]				):





         """simple docstring"""



         pass


      def        a			(self	:				str				):





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


      def        a			(self	:				int				):





         """simple docstring"""



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


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



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


      def        a			(self	:				List[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(a__				)
            __snake_case								=       self._prepare_for_class(a__    ,      a__				)

            __snake_case								=       model(a__    ,      noise=a__				)

            __snake_case								=       copy.deepcopy(self._prepare_for_class(a__    ,      a__				)				)
            __snake_case								=       model(**a__    ,      noise=a__				)
            __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(a__	:				int				):
            __snake_case								=       {}
            for k, v in inputs_dict.items():
               if tf.is_tensor(a__				):
                  __snake_case								=       v.numpy()
               else:
                  __snake_case								=       np.array(a__				)

            return inputs_np_dict

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

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

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


      def        a			(self	:				str    ,      a__	:				Tuple    ,      a__	:				str    ,      a__	:				Tuple				):





         """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(a__				)

         # 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(a__    ,      a__    ,      a__				)


      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(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__				)
         }

         __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(a__				)
         inputs_dict.update({'''noise''': noise}				)

         for main_layer_class in tf_main_layer_classes:
            __snake_case								=       main_layer_class(a__				)

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

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

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


      @slow
      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(a__				)
            __snake_case								=       self._prepare_for_class(a__    ,      a__				)
            __snake_case								=       model(a__    ,      noise=a__				)

            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(a__    ,      saved_model=a__				)
               __snake_case								=       model_class.from_pretrained(a__				)
               __snake_case								=       model(a__    ,      noise=a__				)

               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(a__    ,      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(a__				)
            __snake_case								=       self._prepare_for_class(a__    ,      a__				)

            __snake_case								=       model(a__    ,      noise=a__				)
            __snake_case								=       model.get_config()
            # make sure that returned config is jsonifiable, which is required by keras
            json.dumps(a__				)
            __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(a__				)  # Build model
            new_model.set_weights(model.get_weights()				)
            __snake_case								=       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        a			(self	:				str				):





         """simple docstring"""



         pass


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





         """simple docstring"""



         pass


      @slow
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



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






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







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


      @cached_property
      def        a			(self	:				Any				):





         """simple docstring"""



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


      @slow
      def        a			(self	:				Dict				):





         """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=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)
         __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(**a__    ,      noise=a__				)

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

         __snake_case								=       tf.convert_to_tensor(
             [[-0.0_5_4_8, -1.7_0_2_3, -0.9_3_2_5], [0.3_7_2_1, -0.5_6_7_0, -0.2_2_3_3], [0.8_2_3_5, -1.3_8_7_8, -0.3_5_2_4]]				)

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








import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case_           =      logging.get_logger(__name__)

snake_case_           =      {
    'asapp/sew-d-tiny-100k': 'https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json',
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase  ):
      A_						:    str							=   'sew-d'


      def __init__(self	:				str    ,      a__	:				List[Any]=32    ,      a__	:				Optional[int]=768    ,      a__	:				Union[str, Any]=12    ,      a__	:				Tuple=12    ,      a__	:				int=3072    ,      a__	:				List[Any]=2    ,      a__	:				Dict=512    ,      a__	:				Dict=256    ,      a__	:				List[str]=True    ,      a__	:				List[str]=True    ,      a__	:				str=("p2c", "c2p")    ,      a__	:				str="layer_norm"    ,      a__	:				str="gelu_python"    ,      a__	:				List[Any]=0.1    ,      a__	:				Dict=0.1    ,      a__	:				Optional[Any]=0.1    ,      a__	:				Union[str, Any]=0.0    ,      a__	:				Dict=0.1    ,      a__	:				int=0.0_2    ,      a__	:				str=1E-7    ,      a__	:				Union[str, Any]=1E-5    ,      a__	:				Union[str, Any]="group"    ,      a__	:				Optional[Any]="gelu"    ,      a__	:				List[Any]=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)    ,      a__	:				Union[str, Any]=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)    ,      a__	:				Optional[int]=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)    ,      a__	:				List[str]=False    ,      a__	:				Optional[Any]=128    ,      a__	:				Union[str, Any]=16    ,      a__	:				Optional[Any]=True    ,      a__	:				List[Any]=0.0_5    ,      a__	:				Tuple=10    ,      a__	:				str=2    ,      a__	:				int=0.0    ,      a__	:				Optional[Any]=10    ,      a__	:				List[Any]=0    ,      a__	:				Optional[int]="mean"    ,      a__	:				List[str]=False    ,      a__	:				Tuple=False    ,      a__	:				Optional[int]=256    ,      a__	:				str=0    ,      a__	:				List[Any]=1    ,      a__	:				int=2    ,      **a__	:				List[str]    ,      ):





         """simple docstring"""



         super().__init__(**a__    ,      pad_token_id=a__    ,      bos_token_id=a__    ,      eos_token_id=a__				)
         __snake_case								=       hidden_size
         __snake_case								=       feat_extract_norm
         __snake_case								=       feat_extract_activation
         __snake_case								=       list(a__				)
         __snake_case								=       list(a__				)
         __snake_case								=       list(a__				)
         __snake_case								=       conv_bias
         __snake_case								=       num_conv_pos_embeddings
         __snake_case								=       num_conv_pos_embedding_groups
         __snake_case								=       len(self.conv_dim				)
         __snake_case								=       num_hidden_layers
         __snake_case								=       intermediate_size
         __snake_case								=       squeeze_factor
         __snake_case								=       max_position_embeddings
         __snake_case								=       position_buckets
         __snake_case								=       share_att_key
         __snake_case								=       relative_attention
         __snake_case								=       norm_rel_ebd
         __snake_case								=       list(a__				)
         __snake_case								=       hidden_act
         __snake_case								=       num_attention_heads
         __snake_case								=       hidden_dropout
         __snake_case								=       attention_dropout
         __snake_case								=       activation_dropout
         __snake_case								=       feat_proj_dropout
         __snake_case								=       final_dropout
         __snake_case								=       layer_norm_eps
         __snake_case								=       feature_layer_norm_eps
         __snake_case								=       initializer_range
         __snake_case								=       vocab_size

         if (
             (len(self.conv_stride				) != self.num_feat_extract_layers)
             or (len(self.conv_kernel				) != self.num_feat_extract_layers)
             or (len(self.conv_dim				) != self.num_feat_extract_layers)
         ):
            raise ValueError(
                '''Configuration for convolutional layers is incorrect.'''
                '''It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'''
                f"""but is `len(config.conv_dim) = {len(self.conv_dim				)}`, `len(config.conv_stride)"""
                f"""= {len(self.conv_stride				)}`, `len(config.conv_kernel) = {len(self.conv_kernel				)}`."""				)

         # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
         __snake_case								=       apply_spec_augment
         __snake_case								=       mask_time_prob
         __snake_case								=       mask_time_length
         __snake_case								=       mask_time_min_masks
         __snake_case								=       mask_feature_prob
         __snake_case								=       mask_feature_length
         __snake_case								=       mask_feature_min_masks

         # ctc loss
         __snake_case								=       ctc_loss_reduction
         __snake_case								=       ctc_zero_infinity

         # sequence classification
         __snake_case								=       use_weighted_layer_sum
         __snake_case								=       classifier_proj_size


      @property
      def        a			(self	:				List[Any]				):





         """simple docstring"""



         return functools.reduce(operator.mul    ,      self.conv_stride    ,      1				)

