





"""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 							snake_case							:

      '''simple docstring'''



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


        '''simple docstring'''




        __UpperCAmelCase    :						Tuple       =						parent
        __UpperCAmelCase    :						List[str]       =						batch_size
        __UpperCAmelCase    :						Any       =						image_size
        __UpperCAmelCase    :						str       =						patch_size
        __UpperCAmelCase    :						Any       =						num_channels
        __UpperCAmelCase    :						int       =						is_training
        __UpperCAmelCase    :						int       =						use_labels
        __UpperCAmelCase    :						List[Any]       =						hidden_size
        __UpperCAmelCase    :						Any       =						num_hidden_layers
        __UpperCAmelCase    :						Optional[int]       =						num_attention_heads
        __UpperCAmelCase    :						Union[str, Any]       =						intermediate_size
        __UpperCAmelCase    :						Optional[Any]       =						hidden_act
        __UpperCAmelCase    :						Any       =						hidden_dropout_prob
        __UpperCAmelCase    :						Tuple       =						attention_probs_dropout_prob
        __UpperCAmelCase    :						Union[str, Any]       =						type_sequence_label_size
        __UpperCAmelCase    :						str       =						initializer_range
        __UpperCAmelCase    :						str       =						mask_ratio
        __UpperCAmelCase    :						Optional[Any]       =						scope

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

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase    :						str       =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

        __UpperCAmelCase    :						Dict       =						None
        if self.use_labels:
          __UpperCAmelCase    :						int       =						ids_tensor([self.batch_size]						,				self.type_sequence_label_size )

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

        return config, pixel_values, labels

      def       A_					(							self       :      List[str] ):


        '''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=snake_case_						,				initializer_range=self.initializer_range						,				mask_ratio=self.mask_ratio						,				)

      def       A_					(							self       :      Tuple						,				__lowercase       :      Optional[int]						,				__lowercase       :      int						,				__lowercase       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						TFViTMAEModel(config=snake_case_ )
        __UpperCAmelCase    :						Union[str, Any]       =						model(snake_case_						,				training=snake_case_ )
        self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size) )

      def       A_					(							self       :      Union[str, Any]						,				__lowercase       :      Optional[Any]						,				__lowercase       :      Tuple						,				__lowercase       :      Tuple ):


        '''simple docstring'''




        __UpperCAmelCase    :						int       =						TFViTMAEForPreTraining(snake_case_ )
        __UpperCAmelCase    :						List[str]       =						model(snake_case_						,				training=snake_case_ )
        # expected sequence length = num_patches
        __UpperCAmelCase    :						List[Any]       =						(self.image_size // self.patch_size) ** 2
        __UpperCAmelCase    :						Optional[int]       =						self.patch_size**2 * self.num_channels
        self.parent.assertEqual(result.logits.shape						,				(self.batch_size, num_patches, expected_num_channels) )

        # test greyscale images
        __UpperCAmelCase    :						str       =						1
        __UpperCAmelCase    :						Dict       =						TFViTMAEForPreTraining(snake_case_ )

        __UpperCAmelCase    :						Any       =						floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
        __UpperCAmelCase    :						List[Any]       =						model(snake_case_						,				training=snake_case_ )
        __UpperCAmelCase    :						str       =						self.patch_size**2
        self.parent.assertEqual(result.logits.shape						,				(self.batch_size, num_patches, expected_num_channels) )







      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




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





@require_tf
class 							snake_case							(			_a      ,  _a      ,  unittest.TestCase ):

      '''simple docstring'''



      _A     :       int         =       (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
      _A     :       List[str]         =       {"""feature-extraction""": TFViTMAEModel} if is_tf_available() else {}

      _A     :       List[Any]         =       False
      _A     :       str         =       False
      _A     :       Dict         =       False
      _A     :       Dict         =       False

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						TFViTMAEModelTester(self )
        __UpperCAmelCase    :						Tuple       =						ConfigTester(self						,				config_class=snake_case_						,				has_text_modality=snake_case_						,				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       :      Tuple ):


        '''simple docstring'''




        pass

      def       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase       ,     __UpperCAmelCase    :						Dict       =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						Optional[Any]       =						model_class(snake_case_ )
          self.assertIsInstance(model.get_input_embeddings()						,				(tf.keras.layers.Layer) )
          __UpperCAmelCase    :						List[str]       =						model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(snake_case_						,				tf.keras.layers.Layer ) )

      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        __UpperCAmelCase       ,     __UpperCAmelCase    :						Dict       =						self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						Tuple       =						model_class(snake_case_ )
          __UpperCAmelCase    :						Union[str, Any]       =						inspect.signature(model.call )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          __UpperCAmelCase    :						int       =						[*signature.parameters.keys()]

          __UpperCAmelCase    :						List[Any]       =						['''pixel_values''']
          self.assertListEqual(arg_names[:1]						,				snake_case_ )

      def       A_					(							self       :      Any ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[Any]       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*snake_case_ )

      def       A_					(							self       :      str ):


        '''simple docstring'''




        __UpperCAmelCase    :						int       =						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*snake_case_ )

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase       ,     __UpperCAmelCase    :						int       =						self.model_tester.prepare_config_and_inputs_for_common()

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

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						str       =						model_class(snake_case_ )
          __UpperCAmelCase    :						int       =						self._prepare_for_class(snake_case_						,				snake_case_ )

          __UpperCAmelCase    :						Optional[Any]       =						model(snake_case_						,				noise=snake_case_ )

          __UpperCAmelCase    :						Optional[int]       =						copy.deepcopy(self._prepare_for_class(snake_case_						,				snake_case_ ) )
          __UpperCAmelCase    :						int       =						model(**snake_case_						,				noise=snake_case_ )
          __UpperCAmelCase    :						Optional[int]       =						outputs_dict[0].numpy()
          __UpperCAmelCase    :						Tuple       =						outputs_keywords[0].numpy()

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

      def       A_					(							self       :      Dict ):


        '''simple docstring'''




        np.random.seed(2 )

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

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

        def prepare_numpy_arrays(__lowercase       :      Any ):
          __UpperCAmelCase    :						List[Any]       =						{}
          for k, v in inputs_dict.items():
            if tf.is_tensor(snake_case_ ):
              __UpperCAmelCase    :						str       =						v.numpy()
            else:
              __UpperCAmelCase    :						List[str]       =						np.array(snake_case_ )

          return inputs_np_dict

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						Tuple       =						model_class(snake_case_ )

          __UpperCAmelCase    :						Union[str, Any]       =						self._prepare_for_class(snake_case_						,				snake_case_ )
          __UpperCAmelCase    :						int       =						prepare_numpy_arrays(snake_case_ )

          __UpperCAmelCase    :						Any       =						model(snake_case_						,				noise=snake_case_ )
          __UpperCAmelCase    :						int       =						model(**snake_case_						,				noise=snake_case_ )
          self.assert_outputs_same(snake_case_						,				snake_case_ )

      def       A_					(							self       :      List[Any]						,				__lowercase       :      Optional[Any]						,				__lowercase       :      List[Any]						,				__lowercase       :      Optional[int] ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase    :						Dict       =						int((tf_model.config.image_size // tf_model.config.patch_size) ** 2 )
        __UpperCAmelCase    :						str       =						np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
        __UpperCAmelCase    :						List[Any]       =						tf.constant(snake_case_ )

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

        super().check_pt_tf_models(snake_case_						,				snake_case_						,				snake_case_ )

      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase       ,     __UpperCAmelCase    :						Union[str, Any]       =						self.model_tester.prepare_config_and_inputs_for_common()

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

        __UpperCAmelCase    :						int       =						int((config.image_size // config.patch_size) ** 2 )
        __UpperCAmelCase    :						List[Any]       =						np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
        __UpperCAmelCase    :						Dict       =						tf.convert_to_tensor(snake_case_ )
        inputs_dict.update({'''noise''': noise} )

        for main_layer_class in tf_main_layer_classes:
          __UpperCAmelCase    :						List[Any]       =						main_layer_class(snake_case_ )

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

          __UpperCAmelCase    :						Any       =						tf.keras.Model(snake_case_						,				outputs=main_layer(snake_case_ ) )
          __UpperCAmelCase    :						Optional[int]       =						model(snake_case_ )

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

      @slow
      def       A_					(							self       :      Tuple ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase       ,     __UpperCAmelCase    :						Tuple       =						self.model_tester.prepare_config_and_inputs_for_common()

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

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						List[str]       =						model_class(snake_case_ )
          __UpperCAmelCase    :						Optional[int]       =						self._prepare_for_class(snake_case_						,				snake_case_ )
          __UpperCAmelCase    :						Union[str, Any]       =						model(snake_case_						,				noise=snake_case_ )

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

          with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(snake_case_						,				saved_model=snake_case_ )
            __UpperCAmelCase    :						str       =						model_class.from_pretrained(snake_case_ )
            __UpperCAmelCase    :						List[str]       =						model(snake_case_						,				noise=snake_case_ )

            if model_class.__name__ == "TFViTMAEModel":
              __UpperCAmelCase    :						Optional[Any]       =						after_outputs['''last_hidden_state'''].numpy()
              __UpperCAmelCase    :						str       =						0
            else:
              __UpperCAmelCase    :						Union[str, Any]       =						after_outputs['''logits'''].numpy()
              __UpperCAmelCase    :						Optional[int]       =						0

            __UpperCAmelCase    :						Dict       =						np.amax(np.abs(out_a - out_a ) )
            self.assertLessEqual(snake_case_						,				1e-5 )

      def       A_					(							self       :      List[Any] ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase       ,     __UpperCAmelCase    :						Union[str, Any]       =						self.model_tester.prepare_config_and_inputs_for_common()

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

        for model_class in self.all_model_classes:
          __UpperCAmelCase    :						Any       =						model_class(snake_case_ )
          __UpperCAmelCase    :						Optional[Any]       =						self._prepare_for_class(snake_case_						,				snake_case_ )

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

          self.assert_outputs_same(snake_case_						,				snake_case_ )

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


        '''simple docstring'''




        pass

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


        '''simple docstring'''




        pass







      @slow
      def       A_					(							self       :      int ):


        '''simple docstring'''




        __UpperCAmelCase    :						Optional[int]       =						TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224''' )
        self.assertIsNotNone(snake_case_ )





def       lowerCamelCase_       (   )						->Optional[Any]:





  """simple docstring"""


  __UpperCAmelCase    :						Any       =						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       :      int ):


        '''simple docstring'''




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







      @slow
      def       A_					(							self       :      int ):


        '''simple docstring'''




        np.random.seed(2 )

        __UpperCAmelCase    :						Any       =						TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base''' )

        __UpperCAmelCase    :						List[str]       =						self.default_image_processor
        __UpperCAmelCase    :						str       =						prepare_img()
        __UpperCAmelCase    :						List[str]       =						image_processor(images=snake_case_						,				return_tensors='''tf''' )

        # prepare a noise vector that will be also used for testing the TF model
        # (this way we can ensure that the PT and TF models operate on the same inputs)
        __UpperCAmelCase    :						List[str]       =						ViTMAEConfig()
        __UpperCAmelCase    :						Union[str, Any]       =						int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2 )
        __UpperCAmelCase    :						int       =						np.random.uniform(size=(1, num_patches) )

        # forward pass
        __UpperCAmelCase    :						List[str]       =						model(**snake_case_						,				noise=snake_case_ )

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

        __UpperCAmelCase    :						List[str]       =						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]						,				snake_case_						,				atol=1e-4 )






"""simple docstring"""





import argparse

import torch

from transformers import LxmertConfig, LxmertForPreTraining, load_tf_weights_in_lxmert
from transformers.utils import logging


logging.set_verbosity_info()


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  )						->Optional[int]:





  """simple docstring"""


  __UpperCAmelCase    :						Dict       =						LxmertConfig.from_json_file(UpperCAmelCase_  )
  print(f'''Building PyTorch model from configuration: {config}'''  )
  __UpperCAmelCase    :						Tuple       =						LxmertForPreTraining(UpperCAmelCase_  )

  # Load weights from tf checkpoint
  load_tf_weights_in_lxmert(UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  )

  # Save pytorch-model
  print(f'''Save PyTorch model to {pytorch_dump_path}'''  )
  torch.save(model.state_dict()  ,			UpperCAmelCase_  )


if __name__ == "__main__":
  lowercase__				:Optional[int]				   =       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
  )
  parser.add_argument(
      '--config_file',
      default=None,
      type=str,
      required=True,
      help='The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.',
  )
  parser.add_argument(
      '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
  )
  lowercase__				:Union[str, Any]				   =       parser.parse_args()
  convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)