





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 	_UpperCamelCase	(							A_	):







					'''simple docstring'''

					def 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.config_class(**self.inputs_dict		)
						self.parent.assertTrue(hasattr(__lowercase			,       """hidden_sizes"""		)		)
						self.parent.assertTrue(hasattr(__lowercase			,       """neck_hidden_sizes"""		)		)
						self.parent.assertTrue(hasattr(__lowercase			,       """num_attention_heads"""		)		)







class 	_UpperCamelCase	:







					'''simple docstring'''

					def __init__(    self     :	Union[str, Any]			,       __lowercase     :	Dict			,       __lowercase     :	Any=13			,       __lowercase     :	List[str]=32			,       __lowercase     :	List[str]=2			,       __lowercase     :	Optional[int]=3			,       __lowercase     :	Dict=6_40			,       __lowercase     :	Union[str, Any]=4			,       __lowercase     :	int="silu"			,       __lowercase     :	Dict=3			,       __lowercase     :	int=32			,       __lowercase     :	int=0.1			,       __lowercase     :	Optional[int]=0.1			,       __lowercase     :	Optional[Any]=0.1			,       __lowercase     :	Tuple=0.02			,       __lowercase     :	str=True			,       __lowercase     :	str=True			,       __lowercase     :	Optional[int]=10			,       __lowercase     :	List[Any]=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 							SCREAMING_SNAKE_CASE				(    self     :	List[Any]		):






						'''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 							SCREAMING_SNAKE_CASE				(    self     :	Any		):






						'''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 							SCREAMING_SNAKE_CASE				(    self     :	Any			,       __lowercase     :	int			,       __lowercase     :	Dict			,       __lowercase     :	List[str]			,       __lowercase     :	Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  TFMobileViTModel(config=__lowercase		)
						UpperCAmelCase_		    =  model(__lowercase			,       training=__lowercase		)
						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 							SCREAMING_SNAKE_CASE				(    self     :	int			,       __lowercase     :	int			,       __lowercase     :	Optional[Any]			,       __lowercase     :	Dict			,       __lowercase     :	Tuple		):






						'''simple docstring'''







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





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[str]			,       __lowercase     :	Dict			,       __lowercase     :	Any			,       __lowercase     :	Any			,       __lowercase     :	Dict		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.num_labels
						UpperCAmelCase_		    =  TFMobileViTForSemanticSegmentation(__lowercase		)
						UpperCAmelCase_		    =  UpperCAmelCase_		    =  self.image_size // self.output_stride

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

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





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[Any]		):






						'''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 	_UpperCamelCase	(							A_    ,							A_    ,							unittest.TestCase	):







					'''simple docstring'''

					lowerCamelCase  :  Optional[Any]           =							(
					    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
					    if is_tf_available()
					    else ()
					)
					lowerCamelCase  :  Optional[int]           =							(
					    {'feature-extraction': TFMobileViTModel, 'image-classification': TFMobileViTForImageClassification}
					    if is_tf_available()
					    else {}
					)

					lowerCamelCase  :  Union[str, Any]           =							False
					lowerCamelCase  :  Dict           =							False
					lowerCamelCase  :  Union[str, Any]           =							False
					lowerCamelCase  :  List[Any]           =							False
					lowerCamelCase  :  Union[str, Any]           =							False





					def 							SCREAMING_SNAKE_CASE				(    self     :	int		):






						'''simple docstring'''







						UpperCAmelCase_		    =  TFMobileViTModelTester(self		)
						UpperCAmelCase_		    =  TFMobileViTConfigTester(self			,       config_class=__lowercase			,       has_text_modality=__lowercase		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[Any]		):






						'''simple docstring'''







						self.config_tester.run_common_tests()





					@unittest.skip(reason="""MobileViT does not use inputs_embeds"""		)
					def 							SCREAMING_SNAKE_CASE				(    self     :	List[str]		):






						'''simple docstring'''







						pass





					@unittest.skip(reason="""MobileViT does not support input and output embeddings"""		)
					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[int]		):






						'''simple docstring'''







						pass





					@unittest.skip(reason="""MobileViT does not output attentions"""		)
					def 							SCREAMING_SNAKE_CASE				(    self     :	Dict		):






						'''simple docstring'''







						pass





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[int]		):






						'''simple docstring'''







						UpperCAmelCase_    ,						UpperCAmelCase_		    =  self.model_tester.prepare_config_and_inputs_for_common()

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





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[str]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowercase		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	List[Any]		):






						'''simple docstring'''







						def check_hidden_states_output(__lowercase     :	Any			,       __lowercase     :	str			,       __lowercase     :	List[str]		):
							UpperCAmelCase_		    =  model_class(__lowercase		)

							UpperCAmelCase_		    =  model(**self._prepare_for_class(__lowercase			,       __lowercase		)		)

							UpperCAmelCase_		    =  outputs.hidden_states

							UpperCAmelCase_		    =  5
							self.assertEqual(len(__lowercase		)			,       __lowercase		)

							# 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(__lowercase		)		):
								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(__lowercase			,       __lowercase			,       __lowercase		)

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

							check_hidden_states_output(__lowercase			,       __lowercase			,       __lowercase		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*__lowercase		)





					def 							SCREAMING_SNAKE_CASE				(    self     :	Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_semantic_segmentation(*__lowercase		)





					@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 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''simple docstring'''







						super().test_dataset_conversion()





					def 							SCREAMING_SNAKE_CASE				(    self     :	int			,       __lowercase     :	Dict			,       __lowercase     :	List[Any]			,       __lowercase     :	Dict=2e-1			,       __lowercase     :	Any=2e-1		):






						'''simple docstring'''







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





					@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 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''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(__lowercase		)
								if getattr(__lowercase			,       """hf_compute_loss"""			,       __lowercase		):
									super().test_keras_fit()





					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[Any]		):






						'''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(__lowercase		)
							if getattr(__lowercase			,       """hf_compute_loss"""			,       __lowercase		):
								# 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()			,       __lowercase			,       return_labels=__lowercase		)
								UpperCAmelCase_		    =  prepared_for_class[
								    sorted(prepared_for_class.keys() - inputs_dict.keys()			,       reverse=__lowercase		)[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()			,       __lowercase			,       return_labels=__lowercase		)
								UpperCAmelCase_		    =  {"""input_ids""", """pixel_values""", """input_features"""}
								UpperCAmelCase_		    =  possible_input_names.intersection(set(__lowercase		)		).pop()
								UpperCAmelCase_		    =  prepared_for_class.pop(__lowercase		)

								UpperCAmelCase_		    =  model(__lowercase			,       **__lowercase		)[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()			,       __lowercase			,       return_labels=__lowercase		)
								UpperCAmelCase_		    =  {"""input_ids""", """pixel_values""", """input_features"""}
								UpperCAmelCase_		    =  possible_input_names.intersection(set(__lowercase		)		).pop()
								UpperCAmelCase_		    =  prepared_for_class.pop(__lowercase		)
								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(__lowercase		)
										UpperCAmelCase_		    =  model(__lowercase			,       **__lowercase		)[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()			,       __lowercase			,       return_labels=__lowercase		)
								UpperCAmelCase_		    =  model(__lowercase		)[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()			,       __lowercase			,       return_labels=__lowercase		)

								# 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(__lowercase		)
									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(__lowercase		)

								# 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 							SCREAMING_SNAKE_CASE				(    self     :	str		):






						'''simple docstring'''







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




def  A_(      ):
	UpperCAmelCase_		    =  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
	return image







@require_tf
class 	_UpperCamelCase	(							unittest.TestCase	):







					'''simple docstring'''

					@slow
					def 							SCREAMING_SNAKE_CASE				(    self     :	Any		):






						'''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=__lowercase			,       return_tensors="""tf"""		)

						# forward pass
						UpperCAmelCase_		    =  model(**__lowercase			,       training=__lowercase		)

						# verify the logits
						UpperCAmelCase_		    =  tf.TensorShape((1, 10_00)		)
						self.assertEqual(outputs.logits.shape			,       __lowercase		)

						UpperCAmelCase_		    =  tf.constant([-1.9364, -1.2327, -0.4653]		)

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





					@slow
					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[Any]		):






						'''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=__lowercase			,       return_tensors="""tf"""		)

						# forward pass
						UpperCAmelCase_		    =  model(inputs.pixel_values			,       training=__lowercase		)
						UpperCAmelCase_		    =  outputs.logits

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

						UpperCAmelCase_		    =  tf.constant(
						    [
						        [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
						        [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
						        [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
						    ]		)

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









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


UpperCamelCase__ :			Union[str, Any]       	=			{
    """configuration_mobilevitv2""": [
        """MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """MobileViTV2Config""",
        """MobileViTV2OnnxConfig""",
    ],
}


try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__ :			int       	=			[
				    """MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """MobileViTV2ForImageClassification""",
				    """MobileViTV2ForSemanticSegmentation""",
				    """MobileViTV2Model""",
				    """MobileViTV2PreTrainedModel""",
				]

if TYPE_CHECKING:
				from .configuration_mobilevitva import (
				    MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    MobileViTVaConfig,
				    MobileViTVaOnnxConfig,
				)

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_mobilevitva import (
								    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
								    MobileViTVaForImageClassification,
								    MobileViTVaForSemanticSegmentation,
								    MobileViTVaModel,
								    MobileViTVaPreTrainedModel,
								)

else:
				import sys

				UpperCamelCase__ :			List[Any]       	=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


