
'''simple docstring'''



from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor

class    _a				:





       """simple docstring"""




       def __init__(							self       ,							A__       ,							A__=13       ,							A__=30       ,							A__=2       ,							A__=3       ,							A__=True       ,							A__=True       ,							A__=32       ,							A__=2       ,							A__=4       ,							A__=37       ,							A__="gelu"       ,							A__=0.1       ,							A__=0.1       ,							A__=10       ,							A__=0.02       ,							A__=3       ,							A__=0.6       ,							A__=None       ,							)    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				parent
          _SCREAMING_SNAKE_CASE      =				batch_size
          _SCREAMING_SNAKE_CASE      =				image_size
          _SCREAMING_SNAKE_CASE      =				patch_size
          _SCREAMING_SNAKE_CASE      =				num_channels
          _SCREAMING_SNAKE_CASE      =				is_training
          _SCREAMING_SNAKE_CASE      =				use_labels
          _SCREAMING_SNAKE_CASE      =				hidden_size
          _SCREAMING_SNAKE_CASE      =				num_hidden_layers
          _SCREAMING_SNAKE_CASE      =				num_attention_heads
          _SCREAMING_SNAKE_CASE      =				intermediate_size
          _SCREAMING_SNAKE_CASE      =				hidden_act
          _SCREAMING_SNAKE_CASE      =				hidden_dropout_prob
          _SCREAMING_SNAKE_CASE      =				attention_probs_dropout_prob
          _SCREAMING_SNAKE_CASE      =				type_sequence_label_size
          _SCREAMING_SNAKE_CASE      =				initializer_range
          _SCREAMING_SNAKE_CASE      =				mask_ratio
          _SCREAMING_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)
          _SCREAMING_SNAKE_CASE      =				(image_size // patch_size) ** 2
          _SCREAMING_SNAKE_CASE      =				int(math.ceil((1 - mask_ratio) * (num_patches + 1)      )      )






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

          _SCREAMING_SNAKE_CASE      =				self.get_config()

          return config, pixel_values, labels






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          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        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				TFViTMAEModel(config=A__      )
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							training=A__      )
          self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__      )    ->			Union[str, Any]:
          _SCREAMING_SNAKE_CASE      =				TFViTMAEForPreTraining(A__      )
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							training=A__      )
          # expected sequence length = num_patches
          _SCREAMING_SNAKE_CASE      =				(self.image_size // self.patch_size) ** 2
          _SCREAMING_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
          _SCREAMING_SNAKE_CASE      =				1
          _SCREAMING_SNAKE_CASE      =				TFViTMAEForPreTraining(A__      )

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






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.prepare_config_and_inputs()
          ((_SCREAMING_SNAKE_CASE) ,							(_SCREAMING_SNAKE_CASE) ,							(_SCREAMING_SNAKE_CASE))      =				config_and_inputs
          _SCREAMING_SNAKE_CASE      =				{"""pixel_values""": pixel_values}
          return config, inputs_dict






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





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
       SCREAMING_SNAKE_CASE      =  {'feature-extraction': TFViTMAEModel} if is_tf_available() else {}

       SCREAMING_SNAKE_CASE      =  False
       SCREAMING_SNAKE_CASE      =  False
       SCREAMING_SNAKE_CASE      =  False
       SCREAMING_SNAKE_CASE      =  False






       def        UpperCamelCase				(							self      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				TFViTMAEModelTester(self      )
          _SCREAMING_SNAKE_CASE      =				ConfigTester(self       ,							config_class=A__       ,							has_text_modality=A__       ,							hidden_size=37      )






       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          self.config_tester.run_common_tests()






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






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

             _SCREAMING_SNAKE_CASE      =				["""pixel_values"""]
             self.assertListEqual(arg_names[:1]       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*A__      )






       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*A__      )






       def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
          # make the mask reproducible
          np.random.seed(2      )

          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE      =				model_class(A__      )
             _SCREAMING_SNAKE_CASE      =				self._prepare_for_class(A__       ,							A__      )

             _SCREAMING_SNAKE_CASE      =				model(A__       ,							noise=A__      )

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

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






       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          # make the mask reproducible
          np.random.seed(2      )

          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

          def prepare_numpy_arrays(A__      ):
             _SCREAMING_SNAKE_CASE      =				{}
             for k, v in inputs_dict.items():
                if tf.is_tensor(A__      ):
                   _SCREAMING_SNAKE_CASE      =				v.numpy()
                else:
                   _SCREAMING_SNAKE_CASE      =				np.array(A__      )

             return inputs_np_dict

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

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

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






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__      )    ->			Any:
          # make masks reproducible
          np.random.seed(2      )

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

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

          super().check_pt_tf_models(A__       ,							A__       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			Union[str, Any]:
          # make mask reproducible
          np.random.seed(2      )

          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

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

          for main_layer_class in tf_main_layer_classes:
             _SCREAMING_SNAKE_CASE      =				main_layer_class(A__      )

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

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

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






       @slow
       def        UpperCamelCase				(							self      )    ->			int:
          # make mask reproducible
          np.random.seed(2      )

          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE      =				model_class(A__      )
             _SCREAMING_SNAKE_CASE      =				self._prepare_for_class(A__       ,							A__      )
             _SCREAMING_SNAKE_CASE      =				model(A__       ,							noise=A__      )

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

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

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

                _SCREAMING_SNAKE_CASE      =				np.amax(np.abs(out_a - out_a      )      )
                self.assertLessEqual(A__       ,							1E-5      )






       def        UpperCamelCase				(							self      )    ->			Dict:
          # make mask reproducible
          np.random.seed(2      )

          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_common()

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

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE      =				model_class(A__      )
             _SCREAMING_SNAKE_CASE      =				self._prepare_for_class(A__       ,							A__      )

             _SCREAMING_SNAKE_CASE      =				model(A__       ,							noise=A__      )
             _SCREAMING_SNAKE_CASE      =				model.get_config()
             # make sure that returned config is jsonifiable, which is required by keras
             json.dumps(A__      )
             _SCREAMING_SNAKE_CASE      =				model_class.from_config(model.get_config()      )
             # make sure it also accepts a normal config
             _SCREAMING_SNAKE_CASE      =				model_class.from_config(model.config      )
             _SCREAMING_SNAKE_CASE      =				new_model(A__      )  # Build model
             new_model.set_weights(model.get_weights()      )
             _SCREAMING_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        UpperCamelCase				(							self      )    ->			Any:
          pass






       @unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""      )
       def        UpperCamelCase				(							self      )    ->			Tuple:
          pass






       @slow
       def        UpperCamelCase				(							self      )    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224"""      )
          self.assertIsNotNone(A__      )






def    lowerCAmelCase_	(				)     ->     List[Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""						)
   return image







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





       """simple docstring"""




       @cached_property
       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""      ) if is_vision_available() else None






       @slow
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          # make random mask reproducible across the PT and TF model
          np.random.seed(2      )

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

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

          # forward pass
          _SCREAMING_SNAKE_CASE      =				model(**A__       ,							noise=A__      )

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

          _SCREAMING_SNAKE_CASE      =				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]       ,							A__       ,							atol=1E-4      )








'''simple docstring'''



from pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url

class    _a				(_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  ''
       SCREAMING_SNAKE_CASE      =  'hf-legacy'  # "hf://"" is reserved for hffs






       def __init__(							self       ,							A__ = None       ,							A__ = None       ,							**A__       ,							)    ->			Optional[int]:
          super().__init__(self       ,							**A__      )
          _SCREAMING_SNAKE_CASE      =				repo_info
          _SCREAMING_SNAKE_CASE      =				token
          _SCREAMING_SNAKE_CASE      =				None






       def        UpperCamelCase				(							self      )    ->			Tuple:
          if self.dir_cache is None:
             _SCREAMING_SNAKE_CASE      =				{}
             for hf_file in self.repo_info.siblings:
                # TODO(QL): add sizes
                _SCREAMING_SNAKE_CASE      =				{
                    """name""": hf_file.rfilename,
                    """size""": None,
                    """type""": """file""",
                }
                self.dir_cache.update(
                    {
                        str(A__      ): {"""name""": str(A__      ), """size""": None, """type""": """directory"""}
                        for d in list(PurePosixPath(hf_file.rfilename      ).parents      )[:-1]
                    }      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__ = "rb"       ,							**A__       ,							)    ->			Optional[int]:
          if not isinstance(self.repo_info       ,							A__      ):
             raise NotImplementedError(F"Open is only implemented for dataset repositories, but got {self.repo_info}"      )
          _SCREAMING_SNAKE_CASE      =				hf_hub_url(self.repo_info.id       ,							A__       ,							revision=self.repo_info.sha      )
          return fsspec.open(
              A__       ,							mode=A__       ,							headers=get_authentication_headers_for_url(A__       ,							use_auth_token=self.token      )       ,							client_kwargs={"""trust_env""": True}       ,							).open()






       def        UpperCamelCase				(							self       ,							A__       ,							**A__      )    ->			str:
          self._get_dirs()
          _SCREAMING_SNAKE_CASE      =				self._strip_protocol(A__      )
          if path in self.dir_cache:
             return self.dir_cache[path]
          else:
             raise FileNotFoundError(A__      )






       def        UpperCamelCase				(							self       ,							A__       ,							A__=False       ,							**A__      )    ->			List[Any]:
          self._get_dirs()
          _SCREAMING_SNAKE_CASE      =				PurePosixPath(path.strip("""/"""      )      )
          _SCREAMING_SNAKE_CASE      =				{}
          for p, f in self.dir_cache.items():
             _SCREAMING_SNAKE_CASE      =				PurePosixPath(p.strip("""/"""      )      )
             _SCREAMING_SNAKE_CASE      =				p.parent
             if root == path:
                _SCREAMING_SNAKE_CASE      =				f
          _SCREAMING_SNAKE_CASE      =				list(paths.values()      )
          if detail:
             return out
          else:
             return sorted(f["""name"""] for f in out      )






