
from sklearn.metrics import mean_squared_error

import datasets


_UpperCAmelCase :			Optional[Any]              =	"\\n@article{scikit-learn,\n title={Scikit-learn: Machine Learning in {P}ython},\n author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n journal={Journal of Machine Learning Research},\n volume={12},\n pages={2825--2830},\n year={2011}\n}\n"

_UpperCAmelCase :			Optional[Any]              =	"\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n"


_UpperCAmelCase :			Tuple              =	"\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        \"raw_values\" : Returns a full set of errors in case of multioutput input.\n\n        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric(\"mse\")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {'mse': 0.6123724356957945}\n\n    If you're using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {'mse': array([0.41666667, 1. ])}\n"


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,       _KWARGS_DESCRIPTION       )
class     lowercase (		datasets.Metric       ):
		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				return datasets.MetricInfo(
				    description=_DESCRIPTION    ,				citation=_CITATION    ,				inputs_description=_KWARGS_DESCRIPTION    ,				features=datasets.Features(self._get_feature_types()	)    ,				reference_urls=[
				        'https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html'
				    ]    ,				)
		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				if self.config_name == "multilist":
						return {
						    "predictions": datasets.Sequence(datasets.Value('float'	)	),
						    "references": datasets.Sequence(datasets.Value('float'	)	),
						}
				else:
						return {
						    "predictions": datasets.Value('float'	),
						    "references": datasets.Value('float'	),
						}



		def __UpperCamelCase			(			self    ,				A_    ,				A_    ,				A_=None    ,				A_="uniform_average"    ,				A_=True	)  -> str:
				"""simple docstring"""
				UpperCamelCase          =					mean_squared_error(
				    A_    ,				A_    ,				sample_weight=A_    ,				multioutput=A_    ,				squared=A_	)

				return {"mse": mse}


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



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

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



		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				UpperCamelCase          =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

				UpperCamelCase          =					self.get_config()

				return config, pixel_values, labels



		def __UpperCamelCase			(			self	)  -> Any:
				"""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=A_    ,				initializer_range=self.initializer_range    ,				mask_ratio=self.mask_ratio    ,				)



		def __UpperCamelCase			(			self    ,				A_    ,				A_    ,				A_	)  -> Dict:
				"""simple docstring"""
				UpperCamelCase          =					TFViTMAEModel(config=A_	)
				UpperCamelCase          =					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_	)  -> Tuple:
				"""simple docstring"""
				UpperCamelCase          =					TFViTMAEForPreTraining(A_	)
				UpperCamelCase          =					model(A_    ,				training=A_	)
				# expected sequence length = num_patches
				UpperCamelCase          =					(self.image_size // self.patch_size) ** 2
				UpperCamelCase          =					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          =					1
				UpperCamelCase          =					TFViTMAEForPreTraining(A_	)

				UpperCamelCase          =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
				UpperCamelCase          =					model(A_    ,				training=A_	)
				UpperCamelCase          =					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          =					self.prepare_config_and_inputs()
				((UpperCamelCase)     ,  (UpperCamelCase)     ,  (UpperCamelCase))          =					config_and_inputs
				UpperCamelCase          =					{'pixel_values': pixel_values}
				return config, inputs_dict






@require_tf
class     lowercase (		_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       unittest.TestCase       ):
		__lowercase	:	Union[str, Any]     					=    (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
		__lowercase	:	Optional[int]     					=    {"feature-extraction": TFViTMAEModel} if is_tf_available() else {}

		__lowercase	:	Optional[Any]     					=    False
		__lowercase	:	Optional[Any]     					=    False
		__lowercase	:	Dict     					=    False
		__lowercase	:	List[str]     					=    False



		def __UpperCamelCase			(			self	)  -> Optional[int]:
				"""simple docstring"""
				UpperCamelCase          =					TFViTMAEModelTester(self	)
				UpperCamelCase          =					ConfigTester(self    ,				config_class=A_    ,				has_text_modality=A_    ,				hidden_size=37	)



		def __UpperCamelCase			(			self	)  -> Dict:
				"""simple docstring"""
				self.config_tester.run_common_tests()



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



		def __UpperCamelCase			(			self	)  -> Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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



		def __UpperCamelCase			(			self	)  -> Optional[Any]:
				"""simple docstring"""
				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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

						UpperCamelCase          =					['pixel_values']
						self.assertListEqual(arg_names[:1]    ,				A_	)



		def __UpperCamelCase			(			self	)  -> str:
				"""simple docstring"""
				UpperCamelCase          =					self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*A_	)



		def __UpperCamelCase			(			self	)  -> str:
				"""simple docstring"""
				UpperCamelCase          =					self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_pretraining(*A_	)



		def __UpperCamelCase			(			self	)  -> List[str]:
				"""simple docstring"""
				# make the mask reproducible
				np.random.seed(2	)

				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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

				for model_class in self.all_model_classes:
						UpperCamelCase          =					model_class(A_	)
						UpperCamelCase          =					self._prepare_for_class(A_    ,				A_	)

						UpperCamelCase          =					model(A_    ,				noise=A_	)

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

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



		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				# make the mask reproducible
				np.random.seed(2	)

				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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

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

						return inputs_np_dict

				for model_class in self.all_model_classes:
						UpperCamelCase          =					model_class(A_	)

						UpperCamelCase          =					self._prepare_for_class(A_    ,				A_	)
						UpperCamelCase          =					prepare_numpy_arrays(A_	)

						UpperCamelCase          =					model(A_    ,				noise=A_	)
						UpperCamelCase          =					model(**A_    ,				noise=A_	)
						self.assert_outputs_same(A_    ,				A_	)



		def __UpperCamelCase			(			self    ,				A_    ,				A_    ,				A_	)  -> List[Any]:
				"""simple docstring"""
				# make masks reproducible
				np.random.seed(2	)

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

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

				super().check_pt_tf_models(A_    ,				A_    ,				A_	)



		def __UpperCamelCase			(			self	)  -> Dict:
				"""simple docstring"""
				# make mask reproducible
				np.random.seed(2	)

				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

				UpperCamelCase          =					{
				    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_	)
				}

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

				for main_layer_class in tf_main_layer_classes:
						UpperCamelCase          =					main_layer_class(A_	)

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

						UpperCamelCase          =					tf.keras.Model(A_    ,				outputs=main_layer(A_	)	)
						UpperCamelCase          =					model(A_	)

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



		@slow
		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				# make mask reproducible
				np.random.seed(2	)

				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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

				for model_class in self.all_model_classes:
						UpperCamelCase          =					model_class(A_	)
						UpperCamelCase          =					self._prepare_for_class(A_    ,				A_	)
						UpperCamelCase          =					model(A_    ,				noise=A_	)

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

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

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

								UpperCamelCase          =					np.amax(np.abs(out_a - out_a	)	)
								self.assertLessEqual(A_    ,				1e-5	)



		def __UpperCamelCase			(			self	)  -> Optional[Any]:
				"""simple docstring"""
				# make mask reproducible
				np.random.seed(2	)

				UpperCamelCase     ,  UpperCamelCase          =					self.model_tester.prepare_config_and_inputs_for_common()

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

				for model_class in self.all_model_classes:
						UpperCamelCase          =					model_class(A_	)
						UpperCamelCase          =					self._prepare_for_class(A_    ,				A_	)

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



		@unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'	)
		def __UpperCamelCase			(			self	)  -> Optional[Any]:
				"""simple docstring"""
				pass



		@slow
		def __UpperCamelCase			(			self	)  -> List[Any]:
				"""simple docstring"""
				UpperCamelCase          =					TFViTMAEModel.from_pretrained('google/vit-base-patch16-224'	)
				self.assertIsNotNone(A_	)






def 				A	( )       ->  Optional[int]:
		'''simple docstring'''

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


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



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



		@slow
		def __UpperCamelCase			(			self	)  -> List[str]:
				"""simple docstring"""
				# make random mask reproducible across the PT and TF model
				np.random.seed(2	)

				UpperCamelCase          =					TFViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'	)

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

				# forward pass
				UpperCamelCase          =					model(**A_    ,				noise=A_	)

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

				UpperCamelCase          =					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	)
