




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        snake_case__				(_UpperCamelCase  ):
				"""simple docstring"""



				def 				__UpperCAmelCase				( self							:	List[str]				)   ->	str:
							a       =     self.config_class(**self.inputs_dict				)
							self.parent.assertTrue(hasattr(__lowerCamelCase   , "hidden_sizes"				)				)
							self.parent.assertTrue(hasattr(__lowerCamelCase   , "neck_hidden_sizes"				)				)
							self.parent.assertTrue(hasattr(__lowerCamelCase   , "num_attention_heads"				)				)




class        snake_case__				:
				"""simple docstring"""



				def __init__( self							:	Tuple   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Dict=13   , __lowerCamelCase							:	int=32   , __lowerCamelCase							:	Any=2   , __lowerCamelCase							:	int=3   , __lowerCamelCase							:	Tuple=6_40   , __lowerCamelCase							:	Tuple=4   , __lowerCamelCase							:	Optional[Any]="silu"   , __lowerCamelCase							:	int=3   , __lowerCamelCase							:	Union[str, Any]=32   , __lowerCamelCase							:	Union[str, Any]=0.1   , __lowerCamelCase							:	int=0.1   , __lowerCamelCase							:	str=0.1   , __lowerCamelCase							:	Any=0.02   , __lowerCamelCase							:	Dict=True   , __lowerCamelCase							:	Tuple=True   , __lowerCamelCase							:	str=10   , __lowerCamelCase							:	List[Any]=None   , )   ->	str:
							a       =     parent
							a       =     batch_size
							a       =     image_size
							a       =     patch_size
							a       =     num_channels
							a       =     last_hidden_size
							a       =     num_attention_heads
							a       =     hidden_act
							a       =     conv_kernel_size
							a       =     output_stride
							a       =     hidden_dropout_prob
							a       =     attention_probs_dropout_prob
							a       =     classifier_dropout_prob
							a       =     use_labels
							a       =     is_training
							a       =     num_labels
							a       =     initializer_range
							a       =     scope





				def 				__UpperCAmelCase				( self							:	Any				)   ->	Tuple:
							a       =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

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

							a       =     self.get_config()

							return config, pixel_values, labels, pixel_labels





				def 				__UpperCAmelCase				( self							:	Dict				)   ->	int:
							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							:	Optional[int]   , __lowerCamelCase							:	int   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	int				)   ->	Dict:
							a       =     TFMobileViTModel(config=__lowerCamelCase				)
							a       =     model(__lowerCamelCase   , training=__lowerCamelCase				)
							a       =     a       =     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							:	Optional[int]   , __lowerCamelCase							:	int   , __lowerCamelCase							:	Optional[int]   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	Optional[Any]				)   ->	Union[str, Any]:
							a       =     self.num_labels
							a       =     TFMobileViTForImageClassification(__lowerCamelCase				)
							a       =     model(__lowerCamelCase   , labels=__lowerCamelCase   , training=__lowerCamelCase				)
							self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_labels)				)





				def 				__UpperCAmelCase				( self							:	Tuple   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	Optional[Any]   , __lowerCamelCase							:	Optional[Any]				)   ->	List[Any]:
							a       =     self.num_labels
							a       =     TFMobileViTForSemanticSegmentation(__lowerCamelCase				)
							a       =     a       =     self.image_size // self.output_stride

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

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





				def 				__UpperCAmelCase				( self							:	Any				)   ->	int:
							a       =     self.prepare_config_and_inputs()
							a       ,						a       ,						a       ,						a       =     config_and_inputs
							a       =     {"pixel_values": pixel_values}
							return config, inputs_dict




@require_tf
class        snake_case__				(_UpperCamelCase       ,			_UpperCamelCase       ,			unittest.TestCase  ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_			:				Any									=   (
				    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
				    if is_tf_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE_			:				Union[str, Any]									=   (
				    {"""feature-extraction""": TFMobileViTModel, """image-classification""": TFMobileViTForImageClassification}
				    if is_tf_available()
				    else {}
				)

				SCREAMING_SNAKE_CASE_			:				Optional[int]									=   False
				SCREAMING_SNAKE_CASE_			:				Optional[int]									=   False
				SCREAMING_SNAKE_CASE_			:				Tuple									=   False
				SCREAMING_SNAKE_CASE_			:				str									=   False
				SCREAMING_SNAKE_CASE_			:				Any									=   False





				def 				__UpperCAmelCase				( self							:	int				)   ->	List[str]:
							a       =     TFMobileViTModelTester(self				)
							a       =     TFMobileViTConfigTester(self   , config_class=__lowerCamelCase   , has_text_modality=__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	Dict				)   ->	Optional[int]:
							self.config_tester.run_common_tests()





				@unittest.skip(reason="MobileViT does not use inputs_embeds"				)
				def 				__UpperCAmelCase				( self							:	str				)   ->	Any:
							pass





				@unittest.skip(reason="MobileViT does not support input and output embeddings"				)
				def 				__UpperCAmelCase				( self							:	str				)   ->	Optional[Any]:
							pass





				@unittest.skip(reason="MobileViT does not output attentions"				)
				def 				__UpperCAmelCase				( self							:	str				)   ->	Any:
							pass





				def 				__UpperCAmelCase				( self							:	Optional[int]				)   ->	str:
							a       ,						a       =     self.model_tester.prepare_config_and_inputs_for_common()

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





				def 				__UpperCAmelCase				( self							:	str				)   ->	List[str]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	List[Any]				)   ->	Dict:
							def check_hidden_states_output(__lowerCamelCase							:	Tuple   , __lowerCamelCase							:	List[Any]   , __lowerCamelCase							:	int				):
										a       =     model_class(__lowerCamelCase				)

										a       =     model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase				)				)

										a       =     outputs.hidden_states

										a       =     5
										self.assertEqual(len(__lowerCamelCase				)   , __lowerCamelCase				)

										# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
										# with the width and height being successively divided by 2.
										a       =     2
										for i in range(len(__lowerCamelCase				)				):
													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				)

							a       ,						a       =     self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										a       =     True
										check_hidden_states_output(__lowerCamelCase   , __lowerCamelCase   , __lowerCamelCase				)

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

										check_hidden_states_output(__lowerCamelCase   , __lowerCamelCase   , __lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	Optional[int]				)   ->	Union[str, Any]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase				)





				def 				__UpperCAmelCase				( self							:	List[str]				)   ->	Optional[int]:
							a       =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*__lowerCamelCase				)





				@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							:	Optional[Any]				)   ->	Any:
							super().test_dataset_conversion()





				def 				__UpperCAmelCase				( self							:	Any   , __lowerCamelCase							:	Any   , __lowerCamelCase							:	Tuple   , __lowerCamelCase							:	Dict=2e-1   , __lowerCamelCase							:	int=2e-1				)   ->	int:
							self.assertTrue(np.allclose(__lowerCamelCase   , __lowerCamelCase   , atol=__lowerCamelCase   , rtol=__lowerCamelCase				)				)





				@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							:	List[str]				)   ->	Optional[Any]:
							a       ,						a       =     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":
													a       =     model_class(__lowerCamelCase				)
													if getattr(__lowerCamelCase   , "hf_compute_loss"   , __lowerCamelCase				):
																super().test_keras_fit()





				def 				__UpperCAmelCase				( self							:	Tuple				)   ->	Tuple:
							a       ,						a       =     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":
													a       =     5

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

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

													a       =     model(__lowerCamelCase   , **__lowerCamelCase				)[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
													a       =     self._prepare_for_class(inputs_dict.copy()   , __lowerCamelCase   , return_labels=__lowerCamelCase				)
													a       =     {"input_ids", "pixel_values", "input_features"}
													a       =     possible_input_names.intersection(set(__lowerCamelCase				)				).pop()
													a       =     prepared_for_class.pop(__lowerCamelCase				)
													if "labels" in prepared_for_class:
																a       =     prepared_for_class["labels"].numpy()
																if len(labels.shape				) > 1 and labels.shape[1] != 1:
																			# labels[0] = -100
																			a       =     tf.convert_to_tensor(__lowerCamelCase				)
																			a       =     model(__lowerCamelCase   , **__lowerCamelCase				)[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
													a       =     self._prepare_for_class(inputs_dict.copy()   , __lowerCamelCase   , return_labels=__lowerCamelCase				)
													a       =     model(__lowerCamelCase				)[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
													a       =     self._prepare_for_class(inputs_dict.copy()   , __lowerCamelCase   , return_labels=__lowerCamelCase				)

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

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

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

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

													a       =     tuple(__lowerCamelCase				)

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

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





				@slow
				def 				__UpperCAmelCase				( self							:	Any				)   ->	Optional[Any]:
							for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										a       =     TFMobileViTModel.from_pretrained(__lowerCamelCase				)
										self.assertIsNotNone(__lowerCamelCase				)




def      __magic_name__							(					):






			'''simple docstring'''







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






@require_tf
class        snake_case__				(unittest.TestCase  ):
				"""simple docstring"""



				@slow
				def 				__UpperCAmelCase				( self							:	Tuple				)   ->	str:
							a       =     TFMobileViTForImageClassification.from_pretrained("apple/mobilevit-xx-small"				)

							a       =     MobileViTImageProcessor.from_pretrained("apple/mobilevit-xx-small"				)
							a       =     prepare_img()
							a       =     image_processor(images=__lowerCamelCase   , return_tensors="tf"				)

							# forward pass
							a       =     model(**__lowerCamelCase   , training=__lowerCamelCase				)

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

							a       =     tf.constant([-1.9_364, -1.2_327, -0.4_653]				)

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





				@slow
				def 				__UpperCAmelCase				( self							:	List[Any]				)   ->	List[str]:
							# `from_pt` will be removed
							a       =     TFMobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small"				)

							a       =     MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small"				)

							a       =     prepare_img()
							a       =     image_processor(images=__lowerCamelCase   , return_tensors="tf"				)

							# forward pass
							a       =     model(inputs.pixel_values   , training=__lowerCamelCase				)
							a       =     outputs.logits

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

							a       =     tf.constant(
							    [
							        [[6.9_713, 6.9_786, 7.2_422], [7.2_893, 7.2_825, 7.4_446], [7.6_580, 7.8_797, 7.9_420]],
							        [[-10.6_869, -10.3_250, -10.3_471], [-10.4_228, -9.9_868, -9.7_132], [-11.0_405, -11.0_221, -10.7_318]],
							        [[-3.3_089, -2.8_539, -2.6_740], [-3.2_706, -2.5_621, -2.5_108], [-3.2_534, -2.6_615, -2.6_651]],
							    ]				)

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












class        snake_case__				:
				"""simple docstring"""



				def __init__( self							:	List[str]   , __lowerCamelCase							:	int				)   ->	None:
							a       =     data
							a       =     42
							a       =     42




def      __magic_name__							(					A						:   Node  ):






			'''simple docstring'''







			a       =     0
			a       =     x




def      __magic_name__							(					A						:   Node,    A						:   Node  ):






			'''simple docstring'''







			a       ,						a       =     find_set(A  ), find_set(A  )
			if x == y:
						return

			elif x.rank > y.rank:
						a       =     x
			else:
						a       =     y
						if x.rank == y.rank:
									y.rank += 1




def      __magic_name__							(					A						:   Node  ):






			'''simple docstring'''







			if x != x.parent:
						a       =     find_set(x.parent  )
			return x.parent




def      __magic_name__							(					A						:   Node  ):






			'''simple docstring'''







			a       =     ({0, 1, 2}, {3, 4, 5})
			for s in sets:
						if node.data in s:
									return s
			a       =     F"""{node.data} is not in {sets}"""
			raise ValueError(A  )




def      __magic_name__							(					):






			'''simple docstring'''







			a       =     [Node(A  ) for i in range(6  )]
			for v in vertex:
						make_set(A  )

			union_set(vertex[0],    vertex[1]  )
			union_set(vertex[1],    vertex[2]  )
			union_set(vertex[3],    vertex[4]  )
			union_set(vertex[3],    vertex[5]  )

			for nodea in vertex:
						for nodea in vertex:
									if find_python_set(A  ).isdisjoint(find_python_set(A  )  ):
												assert find_set(A  ) != find_set(A  )
									else:
												assert find_set(A  ) == find_set(A  )


if __name__ == "__main__":
		test_disjoint_set()






