



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    _lowercase (a_   ):
				'''simple docstring'''







				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


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




class    _lowercase :
				'''simple docstring'''







				def __init__(			self  ,			snake_case__  ,			snake_case__=13  ,			snake_case__=32  ,			snake_case__=2  ,			snake_case__=3  ,			snake_case__=640  ,			snake_case__=4  ,			snake_case__="silu"  ,			snake_case__=3  ,			snake_case__=32  ,			snake_case__=0.1  ,			snake_case__=0.1  ,			snake_case__=0.1  ,			snake_case__=0.02  ,			snake_case__=True  ,			snake_case__=True  ,			snake_case__=10  ,			snake_case__=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    _lowerCamelCase       (			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    _lowerCamelCase       (			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    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									UpperCamelCase_			= TFMobileViTModel(config=snake_case__     )
									UpperCamelCase_			= model(snake_case__  ,			training=snake_case__     )
									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    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


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

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									UpperCamelCase_			= self.num_labels
									UpperCamelCase_			= TFMobileViTForSemanticSegmentation(snake_case__     )
									UpperCamelCase_			= UpperCamelCase_			= self.image_size // self.output_stride

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

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

				def    _lowerCamelCase       (			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    _lowercase (a_   , a_   , unittest.TestCase   ):
				'''simple docstring'''







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

				lowercase__     =    False
				lowercase__     =    False
				lowercase__     =    False
				lowercase__     =    False
				lowercase__     =    False

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= TFMobileViTModelTester(self     )
									UpperCamelCase_			= TFMobileViTConfigTester(self  ,			config_class=snake_case__  ,			has_text_modality=snake_case__     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									self.config_tester.run_common_tests()

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

									'''simple docstring'''


									pass

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

									'''simple docstring'''


									pass

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

									'''simple docstring'''


									pass

				def    _lowerCamelCase       (			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(snake_case__     )
														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]  ,			snake_case__     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case__     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									def check_hidden_states_output(snake_case__  ,			snake_case__  ,			snake_case__     ):
														UpperCamelCase_			= model_class(snake_case__     )

														UpperCamelCase_			= model(**self._prepare_for_class(snake_case__  ,			snake_case__     )     )

														UpperCamelCase_			= outputs.hidden_states

														UpperCamelCase_			= 5
														self.assertEqual(len(snake_case__     )  ,			snake_case__     )

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

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

														check_hidden_states_output(snake_case__  ,			snake_case__  ,			snake_case__     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*snake_case__     )

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_semantic_segmentation(*snake_case__     )

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

									'''simple docstring'''


									super().test_dataset_conversion()

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__=2e-1  ,			snake_case__=2e-1     ):

									'''simple docstring'''


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

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

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

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

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

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

									'''simple docstring'''


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




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






@require_tf
class    _lowercase (unittest.TestCase   ):
				'''simple docstring'''







				@slow
				def    _lowerCamelCase       (			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=snake_case__  ,			return_tensors="tf"     )

									# forward pass
									UpperCamelCase_			= model(**snake_case__  ,			training=snake_case__     )

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

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

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

				@slow
				def    _lowerCamelCase       (			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=snake_case__  ,			return_tensors="tf"     )

									# forward pass
									UpperCamelCase_			= model(inputs.pixel_values  ,			training=snake_case__     )
									UpperCamelCase_			= outputs.logits

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

									UpperCamelCase_			= 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]  ,			snake_case__  ,			rtol=1e-4  ,			atol=1e-4     )





import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


UpperCAmelCase :				Any							=1000_0000_0000
UpperCAmelCase :				Any							=1_0000

UpperCAmelCase			,			UpperCAmelCase :				Tuple							=os.path.split(__file__)
UpperCAmelCase :				Tuple							=os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					UpperCamelCase_			= pa.Table.from_pydict({"col": [0] * chunk_size})
					UpperCamelCase_			= pa.concat_tables([table] * (num_examples // chunk_size))
					return datasets.Dataset(_lowerCAmelCase ,						fingerprint="table_100B")


@dataclass
class    _lowercase :
				'''simple docstring'''







				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= np.random.default_rng(self.seed     )
									UpperCamelCase_			= rng.integers(low=self.low  ,			high=self.high  ,			size=self.size     ).tolist()

				def __iter__(			self     ):

									'''simple docstring'''


									return iter(self._sampled_values     )

				def __len__(			self     ):

									'''simple docstring'''


									return self.size




@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[0]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[-1]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[range(len(_lowerCAmelCase) // 2 ,						len(_lowerCAmelCase) // 2 + 10_24)]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[RandIter(0 ,						len(_lowerCAmelCase) ,						10_24 ,						seed=42)]



def  _lowerCAmelCase     ():
					UpperCamelCase_			= {"num examples": SPEED_TEST_N_EXAMPLES}
					UpperCamelCase_			= (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("generating dataset")
					UpperCamelCase_			= generate_100B_dataset(num_examples=_lowerCAmelCase ,						chunk_size=_lowerCAmelCase)
					print("Functions")
					for func in functions:
										print(func.__name__)
										UpperCamelCase_			= func(_lowerCAmelCase)

					with open(_lowerCAmelCase ,						"wb") as f:
										f.write(json.dumps(_lowerCAmelCase).encode("utf-8"))


if __name__ == "__main__":  # useful to run the profiler
	benchmark_table_100B()
