



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_       (   __SCREAMING_SNAKE_CASE			):







	'''simple docstring'''




	def 				_A			(	self       ):



				'''simple docstring'''


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


class lowerCAmelCase_       :







	'''simple docstring'''




	def __init__(	self   ,       lowerCamelCase   ,       lowerCamelCase=13   ,       lowerCamelCase=32   ,       lowerCamelCase=2   ,       lowerCamelCase=3   ,       lowerCamelCase=640   ,       lowerCamelCase=4   ,       lowerCamelCase="silu"   ,       lowerCamelCase=3   ,       lowerCamelCase=32   ,       lowerCamelCase=0.1   ,       lowerCamelCase=0.1   ,       lowerCamelCase=0.1   ,       lowerCamelCase=0.0_2   ,       lowerCamelCase=True   ,       lowerCamelCase=True   ,       lowerCamelCase=10   ,       lowerCamelCase=None   ,       ):



				'''simple docstring'''


				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 				_A			(	self       ):



				'''simple docstring'''


				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 				_A			(	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 				_A			(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):



				'''simple docstring'''


				a__					=  TFMobileViTModel(config=_A       )
				a__					=  model(_A   ,       training=_A       )
				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 				_A			(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):



				'''simple docstring'''


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

	def 				_A			(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):



				'''simple docstring'''


				a__					=  self.num_labels
				a__					=  TFMobileViTForSemanticSegmentation(_A       )
				a__					=  a__					=  self.image_size // self.output_stride

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

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







	def 				_A			(	self       ):



				'''simple docstring'''


				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 lowerCAmelCase_       (   __SCREAMING_SNAKE_CASE ,__SCREAMING_SNAKE_CASE ,unittest.TestCase			):







	'''simple docstring'''




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

	A_       :     Union[str, Any]									=		False
	A_       :     List[str]									=		False
	A_       :     Union[str, Any]									=		False
	A_       :     Tuple									=		False
	A_       :     Optional[int]									=		False

	def 				_A			(	self       ):



				'''simple docstring'''


				a__					=  TFMobileViTModelTester(self       )
				a__					=  TFMobileViTConfigTester(self   ,       config_class=_A   ,       has_text_modality=_A       )

	def 				_A			(	self       ):



				'''simple docstring'''


				self.config_tester.run_common_tests()

	@unittest.skip(reason="""MobileViT does not use inputs_embeds"""       )
	def 				_A			(	self       ):



				'''simple docstring'''


				pass

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



				'''simple docstring'''


				pass

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



				'''simple docstring'''


				pass

	def 				_A			(	self       ):



				'''simple docstring'''


				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 				_A			(	self       ):



				'''simple docstring'''


				a__					=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*_A       )

	def 				_A			(	self       ):



				'''simple docstring'''


				def check_hidden_states_output(lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase       ):
							a__					=  model_class(_A       )

							a__					=  model(**self._prepare_for_class(_A   ,       _A       )       )

							a__					=  outputs.hidden_states

							a__					=  5
							self.assertEqual(len(_A       )   ,       _A       )

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

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

							check_hidden_states_output(_A   ,       _A   ,       _A       )

	def 				_A			(	self       ):



				'''simple docstring'''


				a__					=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*_A       )

	def 				_A			(	self       ):



				'''simple docstring'''


				a__					=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_semantic_segmentation(*_A       )

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



				'''simple docstring'''


				super().test_dataset_conversion()

	def 				_A			(	self   ,       lowerCamelCase   ,       lowerCamelCase   ,       lowerCamelCase=2e-1   ,       lowerCamelCase=2e-1       ):



				'''simple docstring'''


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

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



				'''simple docstring'''


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

	def 				_A			(	self       ):



				'''simple docstring'''


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

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

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

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



				'''simple docstring'''


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




def    UpperCAmelCase					(  ):



			'''simple docstring'''





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





@require_tf
class lowerCAmelCase_       (   unittest.TestCase			):







	'''simple docstring'''




	@slow
	def 				_A			(	self       ):



				'''simple docstring'''


				a__					=  TFMobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""       )

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

				# forward pass
				a__					=  model(**_A   ,       training=_A       )

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

				a__					=  tf.constant([-1.9_3_6_4, -1.2_3_2_7, -0.4_6_5_3]       )

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







	@slow
	def 				_A			(	self       ):



				'''simple docstring'''


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

				# forward pass
				a__					=  model(inputs.pixel_values   ,       training=_A       )
				a__					=  outputs.logits

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

				a__					=  tf.constant(
				    [
				        [[6.9_7_1_3, 6.9_7_8_6, 7.2_4_2_2], [7.2_8_9_3, 7.2_8_2_5, 7.4_4_4_6], [7.6_5_8_0, 7.8_7_9_7, 7.9_4_2_0]],
				        [[-1_0.6_8_6_9, -1_0.3_2_5_0, -1_0.3_4_7_1], [-1_0.4_2_2_8, -9.9_8_6_8, -9.7_1_3_2], [-1_1.0_4_0_5, -1_1.0_2_2_1, -1_0.7_3_1_8]],
				        [[-3.3_0_8_9, -2.8_5_3_9, -2.6_7_4_0], [-3.2_7_0_6, -2.5_6_2_1, -2.5_1_0_8], [-3.2_5_3_4, -2.6_6_1_5, -2.6_6_5_1]],
				    ]       )

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



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


__UpperCAmelCase   						=     logging.get_logger(__name__)





class 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE	):






       """simple docstring"""


       UpperCAmelCase_				       ="speech-encoder-decoder"
       UpperCAmelCase_				       =True



       def __init__(							self ,     **_A	)				->				Dict:
          super().__init__(**_A	)
          if "encoder" not in kwargs or "decoder" not in kwargs:
             raise ValueError(
                 F'''A configuraton of type {self.model_type} cannot be instantiated because not both `encoder` and'''
                 F''' `decoder` sub-configurations are passed, but only {kwargs}'''	)

          SCREAMING_SNAKE_CASE_	=      kwargs.pop('''encoder'''	)
          SCREAMING_SNAKE_CASE_	=      encoder_config.pop('''model_type'''	)
          SCREAMING_SNAKE_CASE_	=      kwargs.pop('''decoder'''	)
          SCREAMING_SNAKE_CASE_	=      decoder_config.pop('''model_type'''	)

          SCREAMING_SNAKE_CASE_	=      AutoConfig.for_model(_A ,     **_A	)
          SCREAMING_SNAKE_CASE_	=      AutoConfig.for_model(_A ,     **_A	)
          SCREAMING_SNAKE_CASE_	=      True



       @classmethod
       def     _UpperCamelCase	(							cls ,     _A ,     _A ,     **_A	)				->				PretrainedConfig:
          logger.info('''Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''	)
          SCREAMING_SNAKE_CASE_	=      True
          SCREAMING_SNAKE_CASE_	=      True

          return cls(encoder=encoder_config.to_dict() ,     decoder=decoder_config.to_dict() ,     **_A	)




       def     _UpperCamelCase	(							self	)				->				Union[str, Any]:
          SCREAMING_SNAKE_CASE_	=      copy.deepcopy(self.__dict__	)
          SCREAMING_SNAKE_CASE_	=      self.encoder.to_dict()
          SCREAMING_SNAKE_CASE_	=      self.decoder.to_dict()
          SCREAMING_SNAKE_CASE_	=      self.__class__.model_type
          return output


