




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




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




     '''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 							__magic_name__							(   self						:	Union[str, Any]				)				->							Tuple:




     '''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 							__magic_name__							(   self						:	str     ,			snake_case_						:	Union[str, Any]     ,			snake_case_						:	Optional[int]     ,			snake_case_						:	Optional[int]				)				->							Tuple:




     '''simple docstring'''

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




    def 							__magic_name__							(   self						:	Optional[Any]     ,			snake_case_						:	List[str]     ,			snake_case_						:	List[Any]     ,			snake_case_						:	Tuple				)				->							Tuple:




     '''simple docstring'''

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

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




    def 							__magic_name__							(   self						:	Tuple				)				->							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      UpperCAmelCase_  (						A_,       A_,       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 							__magic_name__							(   self						:	Dict				)				->							Tuple:




     '''simple docstring'''

     A__	=    TFViTMAEModelTester(self				)
     A__	=    ConfigTester(self     ,			config_class=snake_case_     ,			has_text_modality=snake_case_     ,			hidden_size=37				)




    def 							__magic_name__							(   self						:	List[str]				)				->							Any:




     '''simple docstring'''

     self.config_tester.run_common_tests()




    @unittest.skip(reason="ViTMAE does not use inputs_embeds"				)
    def 							__magic_name__							(   self						:	Tuple				)				->							str:




     '''simple docstring'''

     pass




    def 							__magic_name__							(   self						:	Tuple				)				->							Dict:




     '''simple docstring'''

     A__,				A__	=    self.model_tester.prepare_config_and_inputs_for_common()

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




    def 							__magic_name__							(   self						:	List[str]				)				->							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(snake_case_				)
      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]     ,			snake_case_				)




    def 							__magic_name__							(   self						:	List[Any]				)				->							List[str]:




     '''simple docstring'''

     A__	=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*snake_case_				)




    def 							__magic_name__							(   self						:	int				)				->							Optional[Any]:




     '''simple docstring'''

     A__	=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*snake_case_				)




    def 							__magic_name__							(   self						:	List[str]				)				->							Any:




     '''simple docstring'''

     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(snake_case_				)
      A__	=    self._prepare_for_class(snake_case_     ,			snake_case_				)

      A__	=    model(snake_case_     ,			noise=snake_case_				)

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

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




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




     '''simple docstring'''

     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(snake_case_						:	Dict				):
      A__	=    {}
      for k, v in inputs_dict.items():
       if tf.is_tensor(snake_case_				):
        A__	=    v.numpy()
       else:
        A__	=    np.array(snake_case_				)

      return inputs_np_dict

     for model_class in self.all_model_classes:
      A__	=    model_class(snake_case_				)

      A__	=    self._prepare_for_class(snake_case_     ,			snake_case_				)
      A__	=    prepare_numpy_arrays(snake_case_				)

      A__	=    model(snake_case_     ,			noise=snake_case_				)
      A__	=    model(**snake_case_     ,			noise=snake_case_				)
      self.assert_outputs_same(snake_case_     ,			snake_case_				)




    def 							__magic_name__							(   self						:	Any     ,			snake_case_						:	Optional[int]     ,			snake_case_						:	str     ,			snake_case_						:	List[str]				)				->							List[str]:




     '''simple docstring'''

     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(snake_case_				)

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




    def 							__magic_name__							(   self						:	Any				)				->							Dict:




     '''simple docstring'''

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

     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(snake_case_				)
     inputs_dict.update({"noise": noise}				)

     for main_layer_class in tf_main_layer_classes:
      A__	=    main_layer_class(snake_case_				)

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

      A__	=    tf.keras.Model(snake_case_     ,			outputs=main_layer(snake_case_				)				)
      A__	=    model(snake_case_				)

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




    @slow
    def 							__magic_name__							(   self						:	str				)				->							str:




     '''simple docstring'''

     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(snake_case_				)
      A__	=    self._prepare_for_class(snake_case_     ,			snake_case_				)
      A__	=    model(snake_case_     ,			noise=snake_case_				)

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

       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(snake_case_     ,			1e-5				)




    def 							__magic_name__							(   self						:	Tuple				)				->							List[Any]:




     '''simple docstring'''

     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(snake_case_				)
      A__	=    self._prepare_for_class(snake_case_     ,			snake_case_				)

      A__	=    model(snake_case_     ,			noise=snake_case_				)
      A__	=    model.get_config()
      # make sure that returned config is jsonifiable, which is required by keras
      json.dumps(snake_case_				)
      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(snake_case_				)  # Build model
      new_model.set_weights(model.get_weights()				)
      A__	=    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 							__magic_name__							(   self						:	Optional[Any]				)				->							List[str]:




     '''simple docstring'''

     pass




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




     '''simple docstring'''

     pass




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




     '''simple docstring'''

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







def   _SCREAMING_SNAKE_CASE       (							)						->		List[Any]:
 A__	=    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
 return image






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




    @cached_property
    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Optional[Any]:




     '''simple docstring'''

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




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




     '''simple docstring'''

     np.random.seed(2				)

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

     A__	=    self.default_image_processor
     A__	=    prepare_img()
     A__	=    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)
     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(**snake_case_     ,			noise=snake_case_				)

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

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











"""simple docstring"""



import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class      UpperCAmelCase_  (						A_							):
    lowercase__						 =				['''image_processor''', '''tokenizer''']
    lowercase__						 =				'''ViTImageProcessor'''
    lowercase__						 =				('''CLIPTokenizer''', '''CLIPTokenizerFast''')




    def __init__(   self						:	Optional[Any]     ,			snake_case_						:	Union[str, Any]=None     ,			snake_case_						:	Dict=None     ,			**snake_case_						:	Optional[Any]				)				->							List[str]:




     '''simple docstring'''

     A__	=    None
     if "feature_extractor" in kwargs:
      warnings.warn(
          "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
          " instead."     ,			snake_case_     ,			)
      A__	=    kwargs.pop("feature_extractor"				)

     A__	=    image_processor if image_processor is not None else feature_extractor
     if image_processor is None:
      raise ValueError("You need to specify an `image_processor`."				)
     if tokenizer is None:
      raise ValueError("You need to specify a `tokenizer`."				)

     super().__init__(snake_case_     ,			snake_case_				)




    def __call__(   self						:	int     ,			snake_case_						:	Union[str, Any]=None     ,			snake_case_						:	int=None     ,			snake_case_						:	Dict=None     ,			snake_case_						:	int=None     ,			**snake_case_						:	List[str]				)				->							Union[str, Any]:




     '''simple docstring'''

     if text is None and visual_prompt is None and images is None:
      raise ValueError("You have to specify either text, visual prompt or images."				)

     if text is not None and visual_prompt is not None:
      raise ValueError("You have to specify exactly one type of prompt. Either text or visual prompt."				)

     if text is not None:
      A__	=    self.tokenizer(snake_case_     ,			return_tensors=snake_case_     ,			**snake_case_				)

     if visual_prompt is not None:
      A__	=    self.image_processor(snake_case_     ,			return_tensors=snake_case_     ,			**snake_case_				)

     if images is not None:
      A__	=    self.image_processor(snake_case_     ,			return_tensors=snake_case_     ,			**snake_case_				)

     if visual_prompt is not None and images is not None:
      A__	=    {
          "pixel_values": image_features.pixel_values,
          "conditional_pixel_values": prompt_features.pixel_values,
      }
      return encoding
     elif text is not None and images is not None:
      A__	=    image_features.pixel_values
      return encoding
     elif text is not None:
      return encoding
     elif visual_prompt is not None:
      A__	=    {
          "conditional_pixel_values": prompt_features.pixel_values,
      }
      return encoding
     else:
      return BatchEncoding(data=dict(**snake_case_				)     ,			tensor_type=snake_case_				)




    def 							__magic_name__							(   self						:	Tuple     ,			*snake_case_						:	List[str]     ,			**snake_case_						:	Optional[int]				)				->							str:




     '''simple docstring'''

     return self.tokenizer.batch_decode(*snake_case_     ,			**snake_case_				)




    def 							__magic_name__							(   self						:	Optional[Any]     ,			*snake_case_						:	Any     ,			**snake_case_						:	List[Any]				)				->							List[str]:




     '''simple docstring'''

     return self.tokenizer.decode(*snake_case_     ,			**snake_case_				)




    @property
    def 							__magic_name__							(   self						:	int				)				->							Optional[int]:




     '''simple docstring'''

     warnings.warn(
         "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."     ,			snake_case_     ,			)
     return self.image_processor_class




    @property
    def 							__magic_name__							(   self						:	int				)				->							int:




     '''simple docstring'''

     warnings.warn(
         "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."     ,			snake_case_     ,			)
     return self.image_processor





