



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






						'''simple docstring'''







						def __init__(	self      :			Tuple	,	__lowerCamelCase      :			Optional[Any]	,	__lowerCamelCase      :			List[str]=13	,	__lowerCamelCase      :			Union[str, Any]=30	,	__lowerCamelCase      :			Optional[int]=2	,	__lowerCamelCase      :			Optional[int]=3	,	__lowerCamelCase      :			Union[str, Any]=True	,	__lowerCamelCase      :			List[Any]=True	,	__lowerCamelCase      :			List[str]=32	,	__lowerCamelCase      :			Optional[Any]=2	,	__lowerCamelCase      :			str=4	,	__lowerCamelCase      :			List[Any]=37	,	__lowerCamelCase      :			Tuple="gelu"	,	__lowerCamelCase      :			Tuple=0.1	,	__lowerCamelCase      :			Optional[Any]=0.1	,	__lowerCamelCase      :			List[str]=10	,	__lowerCamelCase      :			List[Any]=0.02	,	__lowerCamelCase      :			int=3	,	__lowerCamelCase      :			List[str]=0.6	,	__lowerCamelCase      :			Union[str, Any]=None	,	)      ->				Optional[int]:
									A							:					Tuple              =						parent
									A							:					str              =						batch_size
									A							:					Any              =						image_size
									A							:					List[str]              =						patch_size
									A							:					List[Any]              =						num_channels
									A							:					Optional[int]              =						is_training
									A							:					List[str]              =						use_labels
									A							:					int              =						hidden_size
									A							:					Any              =						num_hidden_layers
									A							:					str              =						num_attention_heads
									A							:					Union[str, Any]              =						intermediate_size
									A							:					str              =						hidden_act
									A							:					List[Any]              =						hidden_dropout_prob
									A							:					int              =						attention_probs_dropout_prob
									A							:					Tuple              =						type_sequence_label_size
									A							:					Union[str, Any]              =						initializer_range
									A							:					Optional[int]              =						mask_ratio
									A							:					Union[str, Any]              =						scope

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				Tuple:
									A							:					int              =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

									A							:					Union[str, Any]              =						self.get_config()

									return config, pixel_values, labels







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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[Any]	,	__lowerCamelCase      :			List[Any]	,	__lowerCamelCase      :			Optional[Any]	,	__lowerCamelCase      :			str )      ->				Optional[Any]:
									A							:					str              =						TFViTMAEForPreTraining(__lowerCamelCase )
									A							:					Any              =						model(__lowerCamelCase	,	training=__lowerCamelCase )
									# expected sequence length = num_patches
									A							:					Union[str, Any]              =						(self.image_size // self.patch_size) ** 2
									A							:					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
									A							:					List[Any]              =						1
									A							:					Dict              =						TFViTMAEForPreTraining(__lowerCamelCase )

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				Optional[int]:
									A							:					Tuple              =						self.prepare_config_and_inputs()
									((A)       , (A)       , (A))							:					Optional[Any]              =						config_and_inputs
									A							:					Optional[Any]              =						{"pixel_values": pixel_values}
									return config, inputs_dict







@require_tf
class 	lowerCamelCase_    ( _A						,_A						,unittest.TestCase						):






						'''simple docstring'''
						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 			SCREAMING_SNAKE_CASE__		(	self      :			Union[str, Any] )      ->				List[str]:
									A							:					Union[str, Any]              =						TFViTMAEModelTester(self )
									A							:					Any              =						ConfigTester(self	,	config_class=__lowerCamelCase	,	has_text_modality=__lowerCamelCase	,	hidden_size=37 )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[Any] )      ->				Union[str, Any]:
									self.config_tester.run_common_tests()







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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Any )      ->				Optional[int]:
									A       , A							:					List[str]              =						self.model_tester.prepare_config_and_inputs_for_common()

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				Any:
									A       , A							:					Dict              =						self.model_tester.prepare_config_and_inputs_for_common()

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

												A							:					List[str]              =						["pixel_values"]
												self.assertListEqual(arg_names[:1]	,	__lowerCamelCase )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Any )      ->				Any:
									A							:					List[str]              =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__lowerCamelCase )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Union[str, Any] )      ->				Dict:
									A							:					Union[str, Any]              =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*__lowerCamelCase )







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

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

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

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

												A							:					List[str]              =						model(__lowerCamelCase	,	noise=__lowerCamelCase )

												A							:					Dict              =						copy.deepcopy(self._prepare_for_class(__lowerCamelCase	,	__lowerCamelCase ) )
												A							:					List[str]              =						model(**__lowerCamelCase	,	noise=__lowerCamelCase )
												A							:					str              =						outputs_dict[0].numpy()
												A							:					Dict              =						outputs_keywords[0].numpy()

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







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

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

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

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

												return inputs_np_dict

									for model_class in self.all_model_classes:
												A							:					Any              =						model_class(__lowerCamelCase )

												A							:					Optional[int]              =						self._prepare_for_class(__lowerCamelCase	,	__lowerCamelCase )
												A							:					List[str]              =						prepare_numpy_arrays(__lowerCamelCase )

												A							:					Optional[Any]              =						model(__lowerCamelCase	,	noise=__lowerCamelCase )
												A							:					str              =						model(**__lowerCamelCase	,	noise=__lowerCamelCase )
												self.assert_outputs_same(__lowerCamelCase	,	__lowerCamelCase )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict	,	__lowerCamelCase      :			List[Any]	,	__lowerCamelCase      :			List[str]	,	__lowerCamelCase      :			Optional[Any] )      ->				Optional[int]:
									# make masks reproducible
									np.random.seed(2 )

									A							:					Tuple              =						int((tf_model.config.image_size // tf_model.config.patch_size) ** 2 )
									A							:					Dict              =						np.random.uniform(size=(self.model_tester.batch_size, num_patches) )
									A							:					Dict              =						tf.constant(__lowerCamelCase )

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

									super().check_pt_tf_models(__lowerCamelCase	,	__lowerCamelCase	,	__lowerCamelCase )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict )      ->				List[str]:
									# make mask reproducible
									np.random.seed(2 )

									A       , A							:					Optional[Any]              =						self.model_tester.prepare_config_and_inputs_for_common()

									A							:					Any              =						{
									    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 )
									}

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

									for main_layer_class in tf_main_layer_classes:
												A							:					Dict              =						main_layer_class(__lowerCamelCase )

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

												A							:					Tuple              =						tf.keras.Model(__lowerCamelCase	,	outputs=main_layer(__lowerCamelCase ) )
												A							:					Optional[Any]              =						model(__lowerCamelCase )

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







						@slow
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[Any] )      ->				List[str]:
									# make mask reproducible
									np.random.seed(2 )

									A       , A							:					str              =						self.model_tester.prepare_config_and_inputs_for_common()

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

									for model_class in self.all_model_classes:
												A							:					List[str]              =						model_class(__lowerCamelCase )
												A							:					str              =						self._prepare_for_class(__lowerCamelCase	,	__lowerCamelCase )
												A							:					Tuple              =						model(__lowerCamelCase	,	noise=__lowerCamelCase )

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

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

															if model_class.__name__ == "TFViTMAEModel":
																		A							:					Dict              =						after_outputs["last_hidden_state"].numpy()
																		A							:					List[str]              =						0
															else:
																		A							:					str              =						after_outputs["logits"].numpy()
																		A							:					int              =						0

															A							:					int              =						np.amax(np.abs(out_a - out_a ) )
															self.assertLessEqual(__lowerCamelCase	,	1e-5 )







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

									A       , A							:					Optional[int]              =						self.model_tester.prepare_config_and_inputs_for_common()

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

									for model_class in self.all_model_classes:
												A							:					Optional[int]              =						model_class(__lowerCamelCase )
												A							:					Tuple              =						self._prepare_for_class(__lowerCamelCase	,	__lowerCamelCase )

												A							:					List[str]              =						model(__lowerCamelCase	,	noise=__lowerCamelCase )
												A							:					List[str]              =						model.get_config()
												# make sure that returned config is jsonifiable, which is required by keras
												json.dumps(__lowerCamelCase )
												A							:					List[Any]              =						model_class.from_config(model.get_config() )
												# make sure it also accepts a normal config
												A							:					int              =						model_class.from_config(model.config )
												A							:					Union[str, Any]              =						new_model(__lowerCamelCase )  # Build model
												new_model.set_weights(model.get_weights() )
												A							:					str              =						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\n    to get deterministic results." )
						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				int:
									pass







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







						@slow
						def 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				List[Any]:
									A							:					str              =						TFViTMAEModel.from_pretrained("google/vit-base-patch16-224" )
									self.assertIsNotNone(__lowerCamelCase )







def 	UpperCAmelCase   (						):
			A							:					Dict              =						Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
			return image






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






						'''simple docstring'''







						@cached_property
						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Union[str, Any]:
									return ViTImageProcessor.from_pretrained("facebook/vit-mae-base" ) if is_vision_available() else None







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

									A							:					List[Any]              =						TFViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base" )

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

									# forward pass
									A							:					Union[str, Any]              =						model(**__lowerCamelCase	,	noise=__lowerCamelCase )

									# verify the logits
									A							:					Optional[Any]              =						tf.convert_to_tensor([1, 1_96, 7_68] )
									self.assertEqual(outputs.logits.shape	,	__lowerCamelCase )

									A							:					List[str]              =						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 dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Sequence, Value
from .base import TaskTemplate

@dataclass(frozen=_A						)
class 	lowerCamelCase_    ( _A						):






						'''simple docstring'''
						# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
						a__   =  field(default="question-answering-extractive"						,metadata={"include_in_asdict_even_if_is_default": True}						)
						a__   =  Features({"question": Value("string"						), "context": Value("string"						)}						)
						a__   =  Features(
						    {
						        "answers": Sequence(
						            {
						                "text": Value("string"						),
						                "answer_start": Value("int32"						),
						            }						)
						    }						)
						a__   =  "question"
						a__   =  "context"
						a__   =  "answers"







						@property
						def 			SCREAMING_SNAKE_CASE__		(	self      :			Union[str, Any] )      ->				Dict[str, str]:
									return {self.question_column: "question", self.context_column: "context", self.answers_column: "answers"}