






from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, 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 numpy as np
		import tensorflow as tf

		from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
		from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import MobileViTImageProcessor





class __lowerCamelCase					( snake_case__):





	"""simple docstring"""






	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.config_class(**self.inputs_dict       )
							self.parent.assertTrue(hasattr(UpperCAmelCase			,       'hidden_sizes'       )       )
							self.parent.assertTrue(hasattr(UpperCAmelCase			,       'neck_hidden_sizes'       )       )
							self.parent.assertTrue(hasattr(UpperCAmelCase			,       'num_attention_heads'       )       )





class __lowerCamelCase					:





	"""simple docstring"""






	def __init__(				self			,       UpperCAmelCase			,       UpperCAmelCase=13			,       UpperCAmelCase=32			,       UpperCAmelCase=2			,       UpperCAmelCase=3			,       UpperCAmelCase=640			,       UpperCAmelCase=4			,       UpperCAmelCase="silu"			,       UpperCAmelCase=3			,       UpperCAmelCase=32			,       UpperCAmelCase=0.1			,       UpperCAmelCase=0.1			,       UpperCAmelCase=0.1			,       UpperCAmelCase=0.02			,       UpperCAmelCase=True			,       UpperCAmelCase=True			,       UpperCAmelCase=10			,       UpperCAmelCase=None			,       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     parent
							_UpperCAmelCase							   =     batch_size
							_UpperCAmelCase							   =     image_size
							_UpperCAmelCase							   =     patch_size
							_UpperCAmelCase							   =     num_channels
							_UpperCAmelCase							   =     last_hidden_size
							_UpperCAmelCase							   =     num_attention_heads
							_UpperCAmelCase							   =     hidden_act
							_UpperCAmelCase							   =     conv_kernel_size
							_UpperCAmelCase							   =     output_stride
							_UpperCAmelCase							   =     hidden_dropout_prob
							_UpperCAmelCase							   =     attention_probs_dropout_prob
							_UpperCAmelCase							   =     classifier_dropout_prob
							_UpperCAmelCase							   =     use_labels
							_UpperCAmelCase							   =     is_training
							_UpperCAmelCase							   =     num_labels
							_UpperCAmelCase							   =     initializer_range
							_UpperCAmelCase							   =     scope

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

							_UpperCAmelCase							   =     None
							_UpperCAmelCase							   =     None
							if self.use_labels:
													_UpperCAmelCase							   =     ids_tensor([self.batch_size]			,       self.num_labels       )
													_UpperCAmelCase							   =     ids_tensor([self.batch_size, self.image_size, self.image_size]			,       self.num_labels       )

							_UpperCAmelCase							   =     self.get_config()

							return config, pixel_values, labels, pixel_labels

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							return MobileViTConfig(
							    image_size=self.image_size			,       patch_size=self.patch_size			,       num_channels=self.num_channels			,       num_attention_heads=self.num_attention_heads			,       hidden_act=self.hidden_act			,       conv_kernel_size=self.conv_kernel_size			,       output_stride=self.output_stride			,       hidden_dropout_prob=self.hidden_dropout_prob			,       attention_probs_dropout_prob=self.attention_probs_dropout_prob			,       classifier_dropout_prob=self.classifier_dropout_prob			,       initializer_range=self.initializer_range			,       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFMobileViTModel(config=UpperCAmelCase       )
							_UpperCAmelCase							   =     model(UpperCAmelCase			,       training=UpperCAmelCase       )
							_UpperCAmelCase							   =     _UpperCAmelCase							   =     self.image_size // self.output_stride
							self.parent.assertEqual(
							    result.last_hidden_state.shape			,       (self.batch_size, self.last_hidden_size, expected_height, expected_width)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_labels
							_UpperCAmelCase							   =     TFMobileViTForImageClassification(UpperCAmelCase       )
							_UpperCAmelCase							   =     model(UpperCAmelCase			,       labels=UpperCAmelCase			,       training=UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.num_labels)       )

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.num_labels
							_UpperCAmelCase							   =     TFMobileViTForSemanticSegmentation(UpperCAmelCase       )
							_UpperCAmelCase							   =     _UpperCAmelCase							   =     self.image_size // self.output_stride

							_UpperCAmelCase							   =     model(UpperCAmelCase			,       training=UpperCAmelCase       )
							self.parent.assertEqual(
							    result.logits.shape			,       (self.batch_size, self.num_labels, expected_height, expected_width)       )

							_UpperCAmelCase							   =     model(UpperCAmelCase			,       labels=UpperCAmelCase			,       training=UpperCAmelCase       )
							self.parent.assertEqual(
							    result.logits.shape			,       (self.batch_size, self.num_labels, expected_height, expected_width)       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.prepare_config_and_inputs()
							_UpperCAmelCase   ,					_UpperCAmelCase   ,					_UpperCAmelCase   ,					_UpperCAmelCase							   =     config_and_inputs
							_UpperCAmelCase							   =     {'pixel_values': pixel_values}
							return config, inputs_dict





@require_tf
class __lowerCamelCase					( snake_case__	,			snake_case__	,			unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								= (
	    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
	    if is_tf_available()
	    else ()
	)
	UpperCamelCase__								= (
	    {"feature-extraction": TFMobileViTModel, "image-classification": TFMobileViTForImageClassification}
	    if is_tf_available()
	    else {}
	)

	UpperCamelCase__								= False
	UpperCamelCase__								= False
	UpperCamelCase__								= False
	UpperCamelCase__								= False
	UpperCamelCase__								= False

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFMobileViTModelTester(self       )
							_UpperCAmelCase							   =     TFMobileViTConfigTester(self			,       config_class=UpperCAmelCase			,       has_text_modality=UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							self.config_tester.run_common_tests()

	@unittest.skip(reason='MobileViT does not use inputs_embeds'       )
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							pass

	@unittest.skip(reason='MobileViT does not support input and output embeddings'       )
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							pass

	@unittest.skip(reason='MobileViT does not output attentions'       )
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							pass

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase   ,					_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													_UpperCAmelCase							   =     model_class(UpperCAmelCase       )
													_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]			,       UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							def check_hidden_states_output(UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       ):
													_UpperCAmelCase							   =     model_class(UpperCAmelCase       )

													_UpperCAmelCase							   =     model(**self._prepare_for_class(UpperCAmelCase			,       UpperCAmelCase       )       )

													_UpperCAmelCase							   =     outputs.hidden_states

													_UpperCAmelCase							   =     5
													self.assertEqual(len(UpperCAmelCase       )			,       UpperCAmelCase       )

													# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
													# with the width and height being successively divided by 2.
													_UpperCAmelCase							   =     2
													for i in range(len(UpperCAmelCase       )       ):
																			self.assertListEqual(
																			    list(hidden_states[i].shape[-2:]       )			,       [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]			,       )
																			divisor *= 2

													self.assertEqual(self.model_tester.output_stride			,       divisor // 2       )

							_UpperCAmelCase   ,					_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													_UpperCAmelCase							   =     True
													check_hidden_states_output(UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       )

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													_UpperCAmelCase							   =     True

													check_hidden_states_output(UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase       )

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase       )

	@unittest.skipIf(
	    not is_tf_available() or len(tf.config.list_physical_devices('GPU'       )       ) == 0			,       reason='TF does not support backprop for grouped convolutions on CPU.'			,       )
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							super().test_dataset_conversion()

	def 			UpperCamelCase       (				self			,       UpperCAmelCase			,       UpperCAmelCase			,       UpperCAmelCase=2e-1			,       UpperCAmelCase=2e-1       ):
							"""simple docstring"""







							self.assertTrue(np.allclose(UpperCAmelCase			,       UpperCAmelCase			,       atol=UpperCAmelCase			,       rtol=UpperCAmelCase       )       )

	@unittest.skipIf(
	    not is_tf_available() or len(tf.config.list_physical_devices('GPU'       )       ) == 0			,       reason='TF does not support backprop for grouped convolutions on CPU.'			,       )
	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase   ,					_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													# Since `TFMobileViTModel` cannot operate with the default `fit()` method.
													if model_class.__name__ != "TFMobileViTModel":
																			_UpperCAmelCase							   =     model_class(UpperCAmelCase       )
																			if getattr(UpperCAmelCase			,       'hf_compute_loss'			,       UpperCAmelCase       ):
																									super().test_keras_fit()

	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase   ,					_UpperCAmelCase							   =     self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
													# set an ignore index to correctly test the masked loss used in
													# `TFMobileViTForSemanticSegmentation`.
													if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
																			_UpperCAmelCase							   =     5

													_UpperCAmelCase							   =     model_class(UpperCAmelCase       )
													if getattr(UpperCAmelCase			,       'hf_compute_loss'			,       UpperCAmelCase       ):
																			# The number of elements in the loss should be the same as the number of elements in the label
																			_UpperCAmelCase							   =     self._prepare_for_class(inputs_dict.copy()			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )
																			_UpperCAmelCase							   =     prepared_for_class[
																			    sorted(prepared_for_class.keys() - inputs_dict.keys()			,       reverse=UpperCAmelCase       )[0]
																			]
																			_UpperCAmelCase							   =     added_label.shape.as_list()[:1]

																			# Test that model correctly compute the loss with kwargs
																			_UpperCAmelCase							   =     self._prepare_for_class(inputs_dict.copy()			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )
																			_UpperCAmelCase							   =     {'input_ids', 'pixel_values', 'input_features'}
																			_UpperCAmelCase							   =     possible_input_names.intersection(set(UpperCAmelCase       )       ).pop()
																			_UpperCAmelCase							   =     prepared_for_class.pop(UpperCAmelCase       )

																			_UpperCAmelCase							   =     model(UpperCAmelCase			,       **UpperCAmelCase       )[0]
																			self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]       )

																			# Test that model correctly compute the loss when we mask some positions
																			_UpperCAmelCase							   =     self._prepare_for_class(inputs_dict.copy()			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )
																			_UpperCAmelCase							   =     {'input_ids', 'pixel_values', 'input_features'}
																			_UpperCAmelCase							   =     possible_input_names.intersection(set(UpperCAmelCase       )       ).pop()
																			_UpperCAmelCase							   =     prepared_for_class.pop(UpperCAmelCase       )
																			if "labels" in prepared_for_class:
																									_UpperCAmelCase							   =     prepared_for_class['labels'].numpy()
																									if len(labels.shape       ) > 1 and labels.shape[1] != 1:
																															# labels[0] = -100
																															_UpperCAmelCase							   =     tf.convert_to_tensor(UpperCAmelCase       )
																															_UpperCAmelCase							   =     model(UpperCAmelCase			,       **UpperCAmelCase       )[0]
																															self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]       )
																															self.assertTrue(not np.any(np.isnan(loss.numpy()       )       )       )

                # Test that model correctly compute the loss with a dict
																			_UpperCAmelCase							   =     self._prepare_for_class(inputs_dict.copy()			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )
																			_UpperCAmelCase							   =     model(UpperCAmelCase       )[0]
																			self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]       )

																			# Test that model correctly compute the loss with a tuple
																			_UpperCAmelCase							   =     self._prepare_for_class(inputs_dict.copy()			,       UpperCAmelCase			,       return_labels=UpperCAmelCase       )

																			# Get keys that were added with the _prepare_for_class function
																			_UpperCAmelCase							   =     prepared_for_class.keys() - inputs_dict.keys()
																			_UpperCAmelCase							   =     inspect.signature(model.call       ).parameters
																			_UpperCAmelCase							   =     list(signature.keys()       )

																			# Create a dictionary holding the location of the tensors in the tuple
																			_UpperCAmelCase							   =     {0: input_name}
																			for label_key in label_keys:
																									_UpperCAmelCase							   =     signature_names.index(UpperCAmelCase       )
																									_UpperCAmelCase							   =     label_key
																			_UpperCAmelCase							   =     sorted(tuple_index_mapping.items()       )
																			# Initialize a list with their default values, update the values and convert to a tuple
																			_UpperCAmelCase							   =     []

																			for name in signature_names:
																									if name != "kwargs":
																															list_input.append(signature[name].default       )

																			for index, value in sorted_tuple_index_mapping:
																									_UpperCAmelCase							   =     prepared_for_class[value]

																			_UpperCAmelCase							   =     tuple(UpperCAmelCase       )

																			# Send to model
																			_UpperCAmelCase							   =     model(tuple_input[:-1]       )[0]

																			self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]       )

	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_UpperCAmelCase							   =     TFMobileViTModel.from_pretrained(UpperCAmelCase       )
													self.assertIsNotNone(UpperCAmelCase       )



def __A       ( )->    Tuple:





						"""simple docstring"""



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





@require_tf
class __lowerCamelCase					( unittest.TestCase):





	"""simple docstring"""






	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFMobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'       )

							_UpperCAmelCase							   =     MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'       )
							_UpperCAmelCase							   =     prepare_img()
							_UpperCAmelCase							   =     image_processor(images=UpperCAmelCase			,       return_tensors='tf'       )

							# forward pass
							_UpperCAmelCase							   =     model(**UpperCAmelCase			,       training=UpperCAmelCase       )

							# verify the logits
							_UpperCAmelCase							   =     tf.TensorShape((1, 1000)       )
							self.assertEqual(outputs.logits.shape			,       UpperCAmelCase       )

							_UpperCAmelCase							   =     tf.constant([-1.93_64, -1.23_27, -0.46_53]       )

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

	@slow
	def 			UpperCamelCase       (				self       ):
							"""simple docstring"""







							_UpperCAmelCase							   =     TFMobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )

							_UpperCAmelCase							   =     MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )

							_UpperCAmelCase							   =     prepare_img()
							_UpperCAmelCase							   =     image_processor(images=UpperCAmelCase			,       return_tensors='tf'       )

							# forward pass
							_UpperCAmelCase							   =     model(inputs.pixel_values			,       training=UpperCAmelCase       )
							_UpperCAmelCase							   =     outputs.logits

							# verify the logits
							_UpperCAmelCase							   =     tf.TensorShape((1, 21, 32, 32)       )
							self.assertEqual(logits.shape			,       UpperCAmelCase       )

							_UpperCAmelCase							   =     tf.constant(
							    [
							        [[6.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
							        [[-10.68_69, -10.32_50, -10.34_71], [-10.42_28, -9.98_68, -9.71_32], [-11.04_05, -11.02_21, -10.73_18]],
							        [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
							    ]       )

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








def __A       ( __lowerCAmelCase = 100					,						__lowerCAmelCase = 100      )->    int:





						"""simple docstring"""



						return max(
						    sum(int(__lowerCAmelCase      ) for x in str(base**power      )      )
						    for base in range(__lowerCAmelCase      )
						    for power in range(__lowerCAmelCase      )      )


# Tests
if __name__ == "__main__":
		import doctest

		doctest.testmod()
