






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







				def __init__(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=1_3							,       __lowerCAmelCase=3_0							,       __lowerCAmelCase=2							,       __lowerCAmelCase=3							,       __lowerCAmelCase=True							,       __lowerCAmelCase=True							,       __lowerCAmelCase=3_2							,       __lowerCAmelCase=2							,       __lowerCAmelCase=4							,       __lowerCAmelCase=3_7							,       __lowerCAmelCase="gelu"							,       __lowerCAmelCase=0.1							,       __lowerCAmelCase=0.1							,       __lowerCAmelCase=1_0							,       __lowerCAmelCase=0.02							,       __lowerCAmelCase=3							,       __lowerCAmelCase=0.6							,       __lowerCAmelCase=None							,       ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   parent
								__magic_name__		:List[str]   =   batch_size
								__magic_name__		:Any   =   image_size
								__magic_name__		:List[str]   =   patch_size
								__magic_name__		:Optional[int]   =   num_channels
								__magic_name__		:int   =   is_training
								__magic_name__		:Tuple   =   use_labels
								__magic_name__		:List[Any]   =   hidden_size
								__magic_name__		:Dict   =   num_hidden_layers
								__magic_name__		:Union[str, Any]   =   num_attention_heads
								__magic_name__		:Any   =   intermediate_size
								__magic_name__		:List[Any]   =   hidden_act
								__magic_name__		:Tuple   =   hidden_dropout_prob
								__magic_name__		:str   =   attention_probs_dropout_prob
								__magic_name__		:int   =   type_sequence_label_size
								__magic_name__		:Dict   =   initializer_range
								__magic_name__		:Optional[Any]   =   mask_ratio
								__magic_name__		:Any   =   scope

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







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

								__magic_name__		:List[str]   =   None
								if self.use_labels:
												__magic_name__		:Tuple   =   ids_tensor([self.batch_size]							,       self.type_sequence_label_size  )

								__magic_name__		:Dict   =   self.get_config()

								return config, pixel_values, labels







				def       A						(    self  ):
								"""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=__lowerCAmelCase							,       initializer_range=self.initializer_range							,       mask_ratio=self.mask_ratio							,       )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]   =   TFViTMAEModel(config=__lowerCAmelCase  )
								__magic_name__		:Union[str, Any]   =   model(__lowerCAmelCase							,       training=__lowerCAmelCase  )
								self.parent.assertEqual(result.last_hidden_state.shape							,       (self.batch_size, self.seq_length, self.hidden_size)  )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Any   =   TFViTMAEForPreTraining(__lowerCAmelCase  )
								__magic_name__		:List[str]   =   model(__lowerCAmelCase							,       training=__lowerCAmelCase  )
								# expected sequence length = num_patches
								__magic_name__		:List[str]   =   (self.image_size // self.patch_size) ** 2
								__magic_name__		:str   =   self.patch_size**2 * self.num_channels
								self.parent.assertEqual(result.logits.shape							,       (self.batch_size, num_patches, expected_num_channels)  )

								# test greyscale images
								__magic_name__		:Optional[Any]   =   1
								__magic_name__		:Union[str, Any]   =   TFViTMAEForPreTraining(__lowerCAmelCase  )

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







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   self.prepare_config_and_inputs()
								((__magic_name__)  ,					(__magic_name__)  ,					(__magic_name__))		:Tuple   =   config_and_inputs
								__magic_name__		:Optional[int]   =   {'''pixel_values''': pixel_values}
								return config, inputs_dict







@require_tf
class     lowerCamelCase_					(   lowerCamelCase      ,					lowerCamelCase      ,					unittest.TestCase			):
				a__       							=	(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
				a__       							=	{'''feature-extraction''': TFViTMAEModel} if is_tf_available() else {}

				a__       							=	False
				a__       							=	False
				a__       							=	False
				a__       							=	False







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   TFViTMAEModelTester(self  )
								__magic_name__		:Optional[Any]   =   ConfigTester(self							,       config_class=__lowerCAmelCase							,       has_text_modality=__lowerCAmelCase							,       hidden_size=3_7  )







				def       A						(    self  ):
								"""simple docstring"""
								self.config_tester.run_common_tests()







				@unittest.skip(reason='''ViTMAE does not use inputs_embeds'''  )
				def       A						(    self  ):
								"""simple docstring"""
								pass







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Optional[int]   =   self.model_tester.prepare_config_and_inputs_for_common()

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







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Optional[Any]   =   self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												__magic_name__		:Optional[int]   =   model_class(__lowerCAmelCase  )
												__magic_name__		:Dict   =   inspect.signature(model.call  )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												__magic_name__		:Optional[Any]   =   [*signature.parameters.keys()]

												__magic_name__		:Union[str, Any]   =   ['''pixel_values''']
												self.assertListEqual(arg_names[:1]							,       __lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								# make the mask reproducible
								np.random.seed(2  )

								__magic_name__  ,					__magic_name__		:List[Any]   =   self.model_tester.prepare_config_and_inputs_for_common()

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

								for model_class in self.all_model_classes:
												__magic_name__		:int   =   model_class(__lowerCAmelCase  )
												__magic_name__		:str   =   self._prepare_for_class(__lowerCAmelCase							,       __lowerCAmelCase  )

												__magic_name__		:str   =   model(__lowerCAmelCase							,       noise=__lowerCAmelCase  )

												__magic_name__		:str   =   copy.deepcopy(self._prepare_for_class(__lowerCAmelCase							,       __lowerCAmelCase  )  )
												__magic_name__		:Optional[Any]   =   model(**__lowerCAmelCase							,       noise=__lowerCAmelCase  )
												__magic_name__		:Any   =   outputs_dict[0].numpy()
												__magic_name__		:Any   =   outputs_keywords[0].numpy()

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







				def       A						(    self  ):
								"""simple docstring"""
								# make the mask reproducible
								np.random.seed(2  )

								__magic_name__  ,					__magic_name__		:Any   =   self.model_tester.prepare_config_and_inputs_for_common()

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

								def prepare_numpy_arrays(__lowerCAmelCase  ):
												__magic_name__		:Dict   =   {}
												for k, v in inputs_dict.items():
																if tf.is_tensor(__lowerCAmelCase  ):
																				__magic_name__		:Optional[int]   =   v.numpy()
																else:
																				__magic_name__		:Union[str, Any]   =   np.array(__lowerCAmelCase  )

												return inputs_np_dict

								for model_class in self.all_model_classes:
												__magic_name__		:Union[str, Any]   =   model_class(__lowerCAmelCase  )

												__magic_name__		:Optional[Any]   =   self._prepare_for_class(__lowerCAmelCase							,       __lowerCAmelCase  )
												__magic_name__		:Any   =   prepare_numpy_arrays(__lowerCAmelCase  )

												__magic_name__		:Dict   =   model(__lowerCAmelCase							,       noise=__lowerCAmelCase  )
												__magic_name__		:Dict   =   model(**__lowerCAmelCase							,       noise=__lowerCAmelCase  )
												self.assert_outputs_same(__lowerCAmelCase							,       __lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								# make masks reproducible
								np.random.seed(2  )

								__magic_name__		:Tuple   =   int((tf_model.config.image_size // tf_model.config.patch_size) ** 2  )
								__magic_name__		:Optional[int]   =   np.random.uniform(size=(self.model_tester.batch_size, num_patches)  )
								__magic_name__		:List[str]   =   tf.constant(__lowerCAmelCase  )

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

								super().check_pt_tf_models(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								# make mask reproducible
								np.random.seed(2  )

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

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

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

								for main_layer_class in tf_main_layer_classes:
												__magic_name__		:str   =   main_layer_class(__lowerCAmelCase  )

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

												__magic_name__		:List[str]   =   tf.keras.Model(__lowerCAmelCase							,       outputs=main_layer(__lowerCAmelCase  )  )
												__magic_name__		:Tuple   =   model(__lowerCAmelCase  )

												with tempfile.TemporaryDirectory() as tmpdirname:
																__magic_name__		:List[Any]   =   os.path.join(__lowerCAmelCase							,       '''keras_model.h5'''  )
																model.save(__lowerCAmelCase  )
																__magic_name__		:Dict   =   tf.keras.models.load_model(
																    __lowerCAmelCase							,       custom_objects={main_layer_class.__name__: main_layer_class}  )
																assert isinstance(__lowerCAmelCase							,       tf.keras.Model  )
																__magic_name__		:int   =   model(__lowerCAmelCase  )
																self.assert_outputs_same(__lowerCAmelCase							,       __lowerCAmelCase  )







				@slow
				def       A						(    self  ):
								"""simple docstring"""
								# make mask reproducible
								np.random.seed(2  )

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

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

								for model_class in self.all_model_classes:
												__magic_name__		:Tuple   =   model_class(__lowerCAmelCase  )
												__magic_name__		:str   =   self._prepare_for_class(__lowerCAmelCase							,       __lowerCAmelCase  )
												__magic_name__		:List[Any]   =   model(__lowerCAmelCase							,       noise=__lowerCAmelCase  )

												if model_class.__name__ == "TFViTMAEModel":
																__magic_name__		:Dict   =   outputs.last_hidden_state.numpy()
																__magic_name__		:Any   =   0
												else:
																__magic_name__		:Dict   =   outputs.logits.numpy()
																__magic_name__		:Union[str, Any]   =   0

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

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

																__magic_name__		:int   =   np.amax(np.abs(out_a - out_a  )  )
																self.assertLessEqual(__lowerCAmelCase							,       1E-5  )







				def       A						(    self  ):
								"""simple docstring"""
								# make mask reproducible
								np.random.seed(2  )

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

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

								for model_class in self.all_model_classes:
												__magic_name__		:Union[str, Any]   =   model_class(__lowerCAmelCase  )
												__magic_name__		:Union[str, Any]   =   self._prepare_for_class(__lowerCAmelCase							,       __lowerCAmelCase  )

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

												self.assert_outputs_same(__lowerCAmelCase							,       __lowerCAmelCase  )







				@unittest.skip(
				    reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results.'''  )
				def       A						(    self  ):
								"""simple docstring"""
								pass







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







				@slow
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224'''  )
								self.assertIsNotNone(__lowerCAmelCase  )







def   __lowercase       (			):







				"""simple docstring"""






				__magic_name__		:Tuple   =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
				return image

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







				@cached_property
				def       A						(    self  ):
								"""simple docstring"""
								return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base'''  ) if is_vision_available() else None







				@slow
				def       A						(    self  ):
								"""simple docstring"""
								# make random mask reproducible across the PT and TF model
								np.random.seed(2  )

								__magic_name__		:str   =   TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base'''  )

								__magic_name__		:Tuple   =   self.default_image_processor
								__magic_name__		:Dict   =   prepare_img()
								__magic_name__		:Dict   =   image_processor(images=__lowerCAmelCase							,       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)
								__magic_name__		:Tuple   =   ViTMAEConfig()
								__magic_name__		:Union[str, Any]   =   int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2  )
								__magic_name__		:List[str]   =   np.random.uniform(size=(1, num_patches)  )

								# forward pass
								__magic_name__		:str   =   model(**__lowerCAmelCase							,       noise=__lowerCAmelCase  )

								# verify the logits
								__magic_name__		:List[Any]   =   tf.convert_to_tensor([1, 1_9_6, 7_6_8]  )
								self.assertEqual(outputs.logits.shape							,       __lowerCAmelCase  )

								__magic_name__		:int   =   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]							,       __lowerCAmelCase							,       atol=1E-4  )














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     lowerCamelCase_					(   lowerCamelCase			):
				a__       							=	''''''
				a__       							=	'''hf-legacy'''  # "hf://"" is reserved for hffs







				def __init__(    self							,       __lowerCAmelCase = None							,       __lowerCAmelCase = None							,       **__lowerCAmelCase							,       ):
								"""simple docstring"""
								super().__init__(self							,       **__lowerCAmelCase  )
								__magic_name__		:List[Any]   =   repo_info
								__magic_name__		:Dict   =   token
								__magic_name__		:Optional[Any]   =   None







				def       A						(    self  ):
								"""simple docstring"""
								if self.dir_cache is None:
												__magic_name__		:Any   =   {}
												for hf_file in self.repo_info.siblings:
																# TODO(QL): add sizes
																__magic_name__		:Optional[int]   =   {
																    '''name''': hf_file.rfilename,
																    '''size''': None,
																    '''type''': '''file''',
																}
																self.dir_cache.update(
																    {
																        str(__lowerCAmelCase  ): {'''name''': str(__lowerCAmelCase  ), '''size''': None, '''type''': '''directory'''}
																        for d in list(PurePosixPath(hf_file.rfilename  ).parents  )[:-1]
																    }  )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = "rb"							,       **__lowerCAmelCase							,       ):
								"""simple docstring"""
								if not isinstance(self.repo_info							,       __lowerCAmelCase  ):
												raise NotImplementedError(F'''Open is only implemented for dataset repositories, but got {self.repo_info}'''  )
								__magic_name__		:Union[str, Any]   =   hf_hub_url(self.repo_info.id							,       __lowerCAmelCase							,       revision=self.repo_info.sha  )
								return fsspec.open(
								    __lowerCAmelCase							,       mode=__lowerCAmelCase							,       headers=get_authentication_headers_for_url(__lowerCAmelCase							,       use_auth_token=self.token  )							,       client_kwargs={'''trust_env''': True}							,       ).open()







				def       A						(    self							,       __lowerCAmelCase							,       **__lowerCAmelCase  ):
								"""simple docstring"""
								self._get_dirs()
								__magic_name__		:str   =   self._strip_protocol(__lowerCAmelCase  )
								if path in self.dir_cache:
												return self.dir_cache[path]
								else:
												raise FileNotFoundError(__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=False							,       **__lowerCAmelCase  ):
								"""simple docstring"""
								self._get_dirs()
								__magic_name__		:Union[str, Any]   =   PurePosixPath(path.strip('''/'''  )  )
								__magic_name__		:Dict   =   {}
								for p, f in self.dir_cache.items():
												__magic_name__		:int   =   PurePosixPath(p.strip('''/'''  )  )
												__magic_name__		:Tuple   =   p.parent
												if root == path:
																__magic_name__		:Optional[Any]   =   f
								__magic_name__		:List[Any]   =   list(paths.values()  )
								if detail:
												return out
								else:
												return sorted(f['''name'''] for f in out  )






