


from dataclasses import dataclass, field
from typing import ClassVar, Dict

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






@dataclass(frozen=__snake_case				)
class       _UpperCamelCase    (					__snake_case				):

				"""simple docstring"""







				lowerCAmelCase 		= field(default='language-modeling'    ,		metadata={'include_in_asdict_even_if_is_default': True}				)

				lowerCAmelCase 		= Features({'text': Value('string'				)}				)
				lowerCAmelCase 		= Features({}				)
				lowerCAmelCase 		= 'text'

				@property
				def      _UpperCAmelCase					(						self   )	->				Dict[str, str]:
					return {self.text_column: "text"}







from __future__ import annotations

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

import numpy as np

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

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


if is_tf_available():
				import tensorflow as tf

				from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
				from PIL import Image

				from transformers import ViTImageProcessor






class       _UpperCamelCase    :

				"""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				,   )	->				Union[str, Any]:
					A								=				parent
					A								=				batch_size
					A								=				image_size
					A								=				patch_size
					A								=				num_channels
					A								=				is_training
					A								=				use_labels
					A								=				hidden_size
					A								=				num_hidden_layers
					A								=				num_attention_heads
					A								=				intermediate_size
					A								=				hidden_act
					A								=				hidden_dropout_prob
					A								=				attention_probs_dropout_prob
					A								=				type_sequence_label_size
					A								=				initializer_range
					A								=				mask_ratio
					A								=				scope

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

				def      _UpperCAmelCase					(						self   )	->				Optional[Any]:
					A								=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

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

					A								=				self.get_config()

					return config, pixel_values, labels

				def      _UpperCAmelCase					(						self   )	->				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__   )	->				Optional[int]:
					A								=				TFViTMAEModel(config=a__   )
					A								=				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__   )	->				int:
					A								=				TFViTMAEForPreTraining(a__   )
					A								=				model(a__				,   training=a__   )
					# expected sequence length = num_patches
					A								=				(self.image_size // self.patch_size) ** 2
					A								=				self.patch_size**2 * self.num_channels
					self.parent.assertEqual(result.logits.shape				,   (self.batch_size, num_patches, expected_num_channels)   )

					# test greyscale images
					A								=				1
					A								=				TFViTMAEForPreTraining(a__   )

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

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






@require_tf
class       _UpperCamelCase    (					__snake_case    ,		__snake_case    ,		unittest.TestCase				):

				"""simple docstring"""







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

				lowerCAmelCase 		= False
				lowerCAmelCase 		= False
				lowerCAmelCase 		= False
				lowerCAmelCase 		= False

				def      _UpperCAmelCase					(						self   )	->				Dict:
					A								=				TFViTMAEModelTester(self   )
					A								=				ConfigTester(self				,   config_class=a__				,   has_text_modality=a__				,   hidden_size=37   )

				def      _UpperCAmelCase					(						self   )	->				Any:
					self.config_tester.run_common_tests()

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

				def      _UpperCAmelCase					(						self   )	->				int:
					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

				def      _UpperCAmelCase					(						self   )	->				Any:
					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

						A								=				["""pixel_values"""]
						self.assertListEqual(arg_names[:1]				,   a__   )

				def      _UpperCAmelCase					(						self   )	->				Optional[Any]:
					A								=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*a__   )

				def      _UpperCAmelCase					(						self   )	->				int:
					A								=				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_pretraining(*a__   )

				def      _UpperCAmelCase					(						self   )	->				int:
					# make the mask reproducible
					np.random.seed(2   )

					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

					for model_class in self.all_model_classes:
						A								=				model_class(a__   )
						A								=				self._prepare_for_class(a__				,   a__   )

						A								=				model(a__				,   noise=a__   )

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

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

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

					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

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

						return inputs_np_dict

					for model_class in self.all_model_classes:
						A								=				model_class(a__   )

						A								=				self._prepare_for_class(a__				,   a__   )
						A								=				prepare_numpy_arrays(a__   )

						A								=				model(a__				,   noise=a__   )
						A								=				model(**a__				,   noise=a__   )
						self.assert_outputs_same(a__				,   a__   )

				def      _UpperCAmelCase					(						self				,   a__				,   a__				,   a__   )	->				Dict:
					# make masks reproducible
					np.random.seed(2   )

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

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

					super().check_pt_tf_models(a__				,   a__				,   a__   )

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

					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

					A								=				{
					    module_member
					    for model_class in self.all_model_classes
					    for module in (import_module(model_class.__module__   ),)
					    for module_member_name in dir(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__   )
					}

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

					for main_layer_class in tf_main_layer_classes:
						A								=				main_layer_class(a__   )

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

						A								=				tf.keras.Model(a__				,   outputs=main_layer(a__   )   )
						A								=				model(a__   )

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

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

					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

					for model_class in self.all_model_classes:
						A								=				model_class(a__   )
						A								=				self._prepare_for_class(a__				,   a__   )
						A								=				model(a__				,   noise=a__   )

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

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

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

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

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

					A					,			A								=				self.model_tester.prepare_config_and_inputs_for_common()

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

					for model_class in self.all_model_classes:
						A								=				model_class(a__   )
						A								=				self._prepare_for_class(a__				,   a__   )

						A								=				model(a__				,   noise=a__   )
						A								=				model.get_config()
						# make sure that returned config is jsonifiable, which is required by keras
						json.dumps(a__   )
						A								=				model_class.from_config(model.get_config()   )
						# make sure it also accepts a normal config
						A								=				model_class.from_config(model.config   )
						A								=				new_model(a__   )  # Build model
						new_model.set_weights(model.get_weights()   )
						A								=				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   )	->				Tuple:
					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   )	->				str:
					A								=				TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224"""   )
					self.assertIsNotNone(a__   )





def   _lowerCAmelCase				(   ) ->   int:





	"""simple docstring"""







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






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

				"""simple docstring"""







				@cached_property
				def      _UpperCAmelCase					(						self   )	->				Optional[int]:
					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   )

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

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

					# forward pass
					A								=				model(**a__				,   noise=a__   )

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

					A								=				tf.convert_to_tensor(
					    [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]   )

					tf.debugging.assert_near(outputs.logits[0, :3, :3]				,   a__				,   atol=1e-4   )



