
'''simple docstring'''





import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch
      from torch import nn

      from transformers import ViTMAEForPreTraining, ViTMAEModel
      from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import ViTImageProcessor






class 							UpperCAmelCase_							:

 '''simple docstring'''







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



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							parent
   UpperCAmelCase__      :					Optional[int]				       =							batch_size
   UpperCAmelCase__      :					Optional[int]				       =							image_size
   UpperCAmelCase__      :					Tuple				       =							patch_size
   UpperCAmelCase__      :					Dict				       =							num_channels
   UpperCAmelCase__      :					List[Any]				       =							is_training
   UpperCAmelCase__      :					str				       =							use_labels
   UpperCAmelCase__      :					Optional[int]				       =							hidden_size
   UpperCAmelCase__      :					Optional[int]				       =							num_hidden_layers
   UpperCAmelCase__      :					Optional[Any]				       =							num_attention_heads
   UpperCAmelCase__      :					List[Any]				       =							intermediate_size
   UpperCAmelCase__      :					Dict				       =							hidden_act
   UpperCAmelCase__      :					str				       =							hidden_dropout_prob
   UpperCAmelCase__      :					Dict				       =							attention_probs_dropout_prob
   UpperCAmelCase__      :					Dict				       =							type_sequence_label_size
   UpperCAmelCase__      :					Union[str, Any]				       =							initializer_range
   UpperCAmelCase__      :					Dict				       =							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__      :					Union[str, Any]				       =							(image_size // patch_size) ** 2
   UpperCAmelCase__      :					str				       =							int(math.ceil((1 - mask_ratio) * (num_patches + 1) ) )







 def 							UpperCamelCase  (					self		:  str ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[Any]				       =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

   UpperCAmelCase__      :					Optional[Any]				       =							None
   if self.use_labels:
     UpperCAmelCase__      :					Tuple				       =							ids_tensor([self.batch_size]			,				self.type_sequence_label_size )

   UpperCAmelCase__      :					int				       =							self.get_config()

   return config, pixel_values, labels







 def 							UpperCamelCase  (					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			,				hidden_act=self.hidden_act			,				hidden_dropout_prob=self.hidden_dropout_prob			,				attention_probs_dropout_prob=self.attention_probs_dropout_prob			,				is_decoder=_SCREAMING_SNAKE_CASE			,				initializer_range=self.initializer_range			,				mask_ratio=self.mask_ratio			,				)







 def 							UpperCamelCase  (					self		:  Dict			,				snake_case__		:  int			,				snake_case__		:  List[Any]			,				snake_case__		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							ViTMAEModel(config=_SCREAMING_SNAKE_CASE )
   model.to(_SCREAMING_SNAKE_CASE )
   model.eval()
   UpperCAmelCase__      :					Any				       =							model(_SCREAMING_SNAKE_CASE )
   self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size) )







 def 							UpperCamelCase  (					self		:  Optional[Any]			,				snake_case__		:  Optional[int]			,				snake_case__		:  Tuple			,				snake_case__		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							ViTMAEForPreTraining(_SCREAMING_SNAKE_CASE )
   model.to(_SCREAMING_SNAKE_CASE )
   model.eval()
   UpperCAmelCase__      :					Dict				       =							model(_SCREAMING_SNAKE_CASE )
   UpperCAmelCase__      :					List[Any]				       =							(self.image_size // self.patch_size) ** 2
   UpperCAmelCase__      :					Dict				       =							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__      :					Optional[Any]				       =							ViTMAEForPreTraining(_SCREAMING_SNAKE_CASE )
   model.to(_SCREAMING_SNAKE_CASE )
   model.eval()
   UpperCAmelCase__      :					int				       =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
   UpperCAmelCase__      :					Dict				       =							model(_SCREAMING_SNAKE_CASE )
   UpperCAmelCase__      :					Optional[Any]				       =							self.patch_size**2
   self.parent.assertEqual(result.logits.shape			,				(self.batch_size, num_patches, expected_num_channels) )







 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							self.prepare_config_and_inputs()
   UpperCAmelCase__      :					Any				       =							config_and_inputs
   UpperCAmelCase__      :					List[str]				       =							{'''pixel_values''': pixel_values}
   return config, inputs_dict






@require_torch
class 							UpperCAmelCase_							(    _a					,				_a					,				unittest.TestCase				):

 '''simple docstring'''
 lowercase_					:    Tuple				=							(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
 lowercase_					:    List[str]				=							{"feature-extraction": ViTMAEModel} if is_torch_available() else {}

 lowercase_					:    int				=							False
 lowercase_					:    List[Any]				=							False
 lowercase_					:    str				=							False
 lowercase_					:    List[str]				=							False







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[Any]				       =							ViTMAEModelTester(self )
   UpperCAmelCase__      :					List[str]				       =							ConfigTester(self			,				config_class=_SCREAMING_SNAKE_CASE			,				has_text_modality=_SCREAMING_SNAKE_CASE			,				hidden_size=37 )







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   self.config_tester.run_common_tests()







 @unittest.skip(reason="ViTMAE does not use inputs_embeds" )
 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   pass







 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




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

   for model_class in self.all_model_classes:
     UpperCAmelCase__      :					Tuple				       =							model_class(_SCREAMING_SNAKE_CASE )
     self.assertIsInstance(model.get_input_embeddings()			,				(nn.Module) )
     UpperCAmelCase__      :					int				       =							model.get_output_embeddings()
     self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE			,				nn.Linear ) )







 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




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

   for model_class in self.all_model_classes:
     UpperCAmelCase__      :					Any				       =							model_class(_SCREAMING_SNAKE_CASE )
     UpperCAmelCase__      :					Dict				       =							inspect.signature(model.forward )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     UpperCAmelCase__      :					Tuple				       =							[*signature.parameters.keys()]

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







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



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )







 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*_SCREAMING_SNAKE_CASE )







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  Union[str, Any]			,				snake_case__		:  List[str]			,				snake_case__		:  Tuple ):



   '''simple docstring'''




   np.random.seed(2 )

   UpperCAmelCase__      :					List[Any]				       =							int((pt_model.config.image_size // pt_model.config.patch_size) ** 2 )
   UpperCAmelCase__      :					List[Any]				       =							np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
   UpperCAmelCase__      :					Any				       =							torch.from_numpy(_SCREAMING_SNAKE_CASE )

   # Add `noise` argument.
   # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
   UpperCAmelCase__      :					Union[str, Any]				       =							pt_noise

   super().check_pt_tf_models(_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE )







 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     UpperCAmelCase__      :					List[str]				       =							model_class(_SCREAMING_SNAKE_CASE )
     model.to(_SCREAMING_SNAKE_CASE )
     model.eval()
     # make random mask reproducible
     torch.manual_seed(2 )
     with torch.no_grad():
       UpperCAmelCase__      :					Dict				       =							model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE ) )

     UpperCAmelCase__      :					List[str]				       =							outputs[0].cpu().numpy()
     UpperCAmelCase__      :					str				       =							0

     with tempfile.TemporaryDirectory() as tmpdirname:
       model.save_pretrained(_SCREAMING_SNAKE_CASE )
       UpperCAmelCase__      :					Any				       =							model_class.from_pretrained(_SCREAMING_SNAKE_CASE )
       model.to(_SCREAMING_SNAKE_CASE )
       # make random mask reproducible
       torch.manual_seed(2 )
       with torch.no_grad():
         UpperCAmelCase__      :					int				       =							model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE ) )

       # Make sure we don't have nans
       UpperCAmelCase__      :					Tuple				       =							after_outputs[0].cpu().numpy()
       UpperCAmelCase__      :					Optional[int]				       =							0
       UpperCAmelCase__      :					Tuple				       =							np.amax(np.abs(out_a - out_a ) )
       self.assertLessEqual(_SCREAMING_SNAKE_CASE			,				1e-5 )







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



   '''simple docstring'''




   pass







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



   '''simple docstring'''




   pass







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



   '''simple docstring'''




   pass







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



   '''simple docstring'''




   pass







 @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests." )
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   pass







 @slow
 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     UpperCAmelCase__      :					Tuple				       =							ViTMAEModel.from_pretrained(_SCREAMING_SNAKE_CASE )
     self.assertIsNotNone(_SCREAMING_SNAKE_CASE )







def snake_case_     (					):
  UpperCAmelCase__      :					str				       =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
  return image






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

 '''simple docstring'''







 @cached_property
 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




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







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



   '''simple docstring'''




   np.random.seed(2 )

   UpperCAmelCase__      :					List[str]				       =							ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base" ).to(_SCREAMING_SNAKE_CASE )

   UpperCAmelCase__      :					int				       =							self.default_image_processor
   UpperCAmelCase__      :					str				       =							prepare_img()
   UpperCAmelCase__      :					Union[str, Any]				       =							image_processor(images=_SCREAMING_SNAKE_CASE			,				return_tensors="pt" ).to(_SCREAMING_SNAKE_CASE )

   # 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__      :					Any				       =							ViTMAEConfig()
   UpperCAmelCase__      :					Any				       =							int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2 )
   UpperCAmelCase__      :					Optional[Any]				       =							np.random.uniform(size=(1, num_patches) )

   # forward pass
   with torch.no_grad():
     UpperCAmelCase__      :					int				       =							model(**_SCREAMING_SNAKE_CASE			,				noise=torch.from_numpy(_SCREAMING_SNAKE_CASE ).to(device=_SCREAMING_SNAKE_CASE ) )

   # verify the logits
   UpperCAmelCase__      :					Tuple				       =							torch.Size((1, 1_96, 7_68) )
   self.assertEqual(outputs.logits.shape			,				_SCREAMING_SNAKE_CASE )

   UpperCAmelCase__      :					Optional[Any]				       =							torch.tensor(
       [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]] )

   self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,				expected_slice.to(_SCREAMING_SNAKE_CASE )			,				atol=1e-4 ) )









__a							     =					'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     if not isinstance(_lowercase					,		_lowercase ):
          UpperCAmelCase_     :				Union[str, Any]  =     f'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
          raise TypeError(_lowercase )

     UpperCAmelCase_     :				Any  =     ''''''.join(bin(_lowercase )[2:].zfill(8 ) for byte in data )

     UpperCAmelCase_     :				Any  =     len(_lowercase ) % 6 != 0

     if padding_needed:
          # The padding that will be added later
          UpperCAmelCase_     :				Union[str, Any]  =     B'''=''' * ((6 - len(_lowercase ) % 6) // 2)

          # Append binary_stream with arbitrary binary digits (0's by default) to make its
          # length a multiple of 6.
          binary_stream += "0" * (6 - len(_lowercase ) % 6)
     else:
          UpperCAmelCase_     :				int  =     B''''''

     # Encode every 6 binary digits to their corresponding Base64 character
     return (
         "".join(
             B64_CHARSET[int(binary_stream[index : index + 6]					,		2 )]
             for index in range(0					,		len(_lowercase )					,		6 ) ).encode()
         + padding
     )







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     if not isinstance(_lowercase					,		_lowercase ) and not isinstance(_lowercase					,		_lowercase ):
          UpperCAmelCase_     :				Tuple  =     (
              '''argument should be a bytes-like object or ASCII string, '''
              f'''not \'{encoded_data.__class__.__name__}\''''
          )
          raise TypeError(_lowercase )

     # In case encoded_data is a bytes-like object, make sure it contains only
     # ASCII characters so we convert it to a string object
     if isinstance(_lowercase					,		_lowercase ):
          try:
               UpperCAmelCase_     :				Any  =     encoded_data.decode('''utf-8''' )
          except UnicodeDecodeError:
               raise ValueError('''base64 encoded data should only contain ASCII characters''' )

     UpperCAmelCase_     :				str  =     encoded_data.count('''=''' )

     # Check if the encoded string contains non base64 characters
     if padding:
          assert all(
              char in B64_CHARSET for char in encoded_data[:-padding] ), "Invalid base64 character(s) found."
     else:
          assert all(
              char in B64_CHARSET for char in encoded_data ), "Invalid base64 character(s) found."

     # Check the padding
     assert len(_lowercase ) % 4 == 0 and padding < 3, "Incorrect padding"

     if padding:
          # Remove padding if there is one
          UpperCAmelCase_     :				List[Any]  =     encoded_data[:-padding]

          UpperCAmelCase_     :				List[Any]  =     ''''''.join(
              bin(B64_CHARSET.index(_lowercase ) )[2:].zfill(6 ) for char in encoded_data )[: -padding * 2]
     else:
          UpperCAmelCase_     :				Tuple  =     ''''''.join(
              bin(B64_CHARSET.index(_lowercase ) )[2:].zfill(6 ) for char in encoded_data )

     UpperCAmelCase_     :				str  =     [
         int(binary_stream[index : index + 8]					,		2 )
         for index in range(0					,		len(_lowercase )					,		8 )
     ]

     return bytes(_lowercase )


if __name__ == "__main__":
     import doctest

     doctest.testmod()