
import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
					from transformers.models.mobilevitva.modeling_mobilevitva import (
					    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
					    make_divisible,
					)


if is_vision_available():
					from PIL import Image

					from transformers import MobileViTImageProcessor

class    lowerCAmelCase__     (	__lowercase			):






		def  __A       (		self		:							int  )     ->							Union[str, Any]:
						__lowerCamelCase    					=	self.config_class(**self.inputs_dict  )
						self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__							,			'''width_multiplier'''  )  )





class    lowerCAmelCase__     :






		def __init__(		self		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							List[Any]=13							,			SCREAMING_SNAKE_CASE__		:							str=64							,			SCREAMING_SNAKE_CASE__		:							List[str]=2							,			SCREAMING_SNAKE_CASE__		:							Dict=3							,			SCREAMING_SNAKE_CASE__		:							int="swish"							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=3							,			SCREAMING_SNAKE_CASE__		:							List[str]=32							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=0.1							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=0.02							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=True							,			SCREAMING_SNAKE_CASE__		:							Union[str, Any]=True							,			SCREAMING_SNAKE_CASE__		:							str=10							,			SCREAMING_SNAKE_CASE__		:							Tuple=None							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=0.25							,			SCREAMING_SNAKE_CASE__		:							str=0.0							,			SCREAMING_SNAKE_CASE__		:							List[str]=0.0							,			)     ->							Optional[Any]:
						__lowerCamelCase    					=	parent
						__lowerCamelCase    					=	batch_size
						__lowerCamelCase    					=	image_size
						__lowerCamelCase    					=	patch_size
						__lowerCamelCase    					=	num_channels
						__lowerCamelCase    					=	make_divisible(5_12 * width_multiplier							,			divisor=8  )
						__lowerCamelCase    					=	hidden_act
						__lowerCamelCase    					=	conv_kernel_size
						__lowerCamelCase    					=	output_stride
						__lowerCamelCase    					=	classifier_dropout_prob
						__lowerCamelCase    					=	use_labels
						__lowerCamelCase    					=	is_training
						__lowerCamelCase    					=	num_labels
						__lowerCamelCase    					=	initializer_range
						__lowerCamelCase    					=	scope
						__lowerCamelCase    					=	width_multiplier
						__lowerCamelCase    					=	ffn_dropout
						__lowerCamelCase    					=	attn_dropout






		def  __A       (		self		:							Union[str, Any]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

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

						__lowerCamelCase    					=	self.get_config()

						return config, pixel_values, labels, pixel_labels






		def  __A       (		self		:							str  )     ->							int:
						return MobileViTVaConfig(
						    image_size=self.image_size							,			patch_size=self.patch_size							,			num_channels=self.num_channels							,			hidden_act=self.hidden_act							,			conv_kernel_size=self.conv_kernel_size							,			output_stride=self.output_stride							,			classifier_dropout_prob=self.classifier_dropout_prob							,			initializer_range=self.initializer_range							,			width_multiplier=self.width_multiplier							,			ffn_dropout=self.ffn_dropout_prob							,			attn_dropout=self.attn_dropout_prob							,			)






		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							str							,			SCREAMING_SNAKE_CASE__		:							Any							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							Any:
						__lowerCamelCase    					=	MobileViTVaModel(config=SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.eval()
						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__  )
						self.parent.assertEqual(
						    result.last_hidden_state.shape							,			(
						        self.batch_size,
						        self.last_hidden_size,
						        self.image_size // self.output_stride,
						        self.image_size // self.output_stride,
						    )							,			)






		def  __A       (		self		:							int							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							Dict  )     ->							List[Any]:
						__lowerCamelCase    					=	self.num_labels
						__lowerCamelCase    					=	MobileViTVaForImageClassification(SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.eval()
						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__							,			labels=SCREAMING_SNAKE_CASE__  )
						self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.num_labels)  )






		def  __A       (		self		:							Any							,			SCREAMING_SNAKE_CASE__		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							Union[str, Any]  )     ->							Any:
						__lowerCamelCase    					=	self.num_labels
						__lowerCamelCase    					=	MobileViTVaForSemanticSegmentation(SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.eval()
						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__  )
						self.parent.assertEqual(
						    result.logits.shape							,			(
						        self.batch_size,
						        self.num_labels,
						        self.image_size // self.output_stride,
						        self.image_size // self.output_stride,
						    )							,			)
						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__							,			labels=SCREAMING_SNAKE_CASE__  )
						self.parent.assertEqual(
						    result.logits.shape							,			(
						        self.batch_size,
						        self.num_labels,
						        self.image_size // self.output_stride,
						        self.image_size // self.output_stride,
						    )							,			)







		def  __A       (		self		:							List[str]  )     ->							Dict:
						__lowerCamelCase    					=	self.prepare_config_and_inputs()
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	config_and_inputs
						__lowerCamelCase    					=	{'''pixel_values''': pixel_values}
						return config, inputs_dict





@require_torch
class    lowerCAmelCase__     (	__lowercase  ,	__lowercase  ,	unittest.TestCase			):
		a__       :						Optional[int]        =		(
		    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
		    if is_torch_available()
		    else ()
		)

		a__       :						Any        =		(
		    {
		        """feature-extraction""": MobileViTVaModel,
		        """image-classification""": MobileViTVaForImageClassification,
		        """image-segmentation""": MobileViTVaForSemanticSegmentation,
		    }
		    if is_torch_available()
		    else {}
		)

		a__       :						List[str]        =		False
		a__       :						Optional[Any]        =		False
		a__       :						str        =		False
		a__       :						Optional[Any]        =		False






		def  __A       (		self		:							List[Any]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	MobileViTVaModelTester(self  )
						__lowerCamelCase    					=	MobileViTVaConfigTester(self							,			config_class=SCREAMING_SNAKE_CASE__							,			has_text_modality=SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							int  )     ->							Tuple:
						self.config_tester.run_common_tests()






		@unittest.skip(reason='''MobileViTV2 does not use inputs_embeds'''  )
		def  __A       (		self		:							str  )     ->							str:
						pass






		@unittest.skip(reason='''MobileViTV2 does not support input and output embeddings'''  )
		def  __A       (		self		:							Tuple  )     ->							Optional[Any]:
						pass






		@unittest.skip(reason='''MobileViTV2 does not output attentions'''  )
		def  __A       (		self		:							Union[str, Any]  )     ->							Tuple:
						pass






		@require_torch_multi_gpu
		@unittest.skip(reason='''Got `CUDA error: misaligned address` for tests after this one being run.'''  )
		def  __A       (		self		:							List[Any]  )     ->							List[Any]:
						pass






		@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''  )
		def  __A       (		self		:							int  )     ->							Dict:
						pass






		def  __A       (		self		:							Optional[Any]  )     ->							Any:
						__lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										__lowerCamelCase    					=	model_class(SCREAMING_SNAKE_CASE__  )
										__lowerCamelCase    					=	inspect.signature(model.forward  )
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										__lowerCamelCase    					=	[*signature.parameters.keys()]

										__lowerCamelCase    					=	['''pixel_values''']
										self.assertListEqual(arg_names[:1]							,			SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							List[str]  )     ->							Any:
						__lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							Dict  )     ->							Optional[Any]:
						def check_hidden_states_output(SCREAMING_SNAKE_CASE__		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[int]  ):
										__lowerCamelCase    					=	model_class(SCREAMING_SNAKE_CASE__  )
										model.to(SCREAMING_SNAKE_CASE__  )
										model.eval()

										with torch.no_grad():
														__lowerCamelCase    					=	model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  )  )

										__lowerCamelCase    					=	outputs.hidden_states

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

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

						__lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs_for_common()

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

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

										check_hidden_states_output(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							List[str]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							int  )     ->							Tuple:
						__lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_SNAKE_CASE__  )







		@slow
		def  __A       (		self		:							Optional[Any]  )     ->							Any:
						for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										__lowerCamelCase    					=	MobileViTVaModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
										self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )





def  __magic_name__ (   )			->   Optional[int]:
				__lowerCamelCase    					=	Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
				return image







@require_torch
@require_vision
class    lowerCAmelCase__     (	unittest.TestCase			):






		@cached_property
		def  __A       (		self		:							str  )     ->							Dict:
						return (
						    MobileViTImageProcessor.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''  )
						    if is_vision_available()
						    else None
						)






		@slow
		def  __A       (		self		:							Optional[int]  )     ->							Union[str, Any]:
						__lowerCamelCase    					=	MobileViTVaForImageClassification.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''  ).to(
						    SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	self.default_image_processor
						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(images=SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )

						# forward pass
						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )

						# verify the logits
						__lowerCamelCase    					=	torch.Size((1, 10_00)  )
						self.assertEqual(outputs.logits.shape							,			SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	torch.tensor([-1.6_336e00, -7.3_204e-02, -5.1_883e-01]  ).to(SCREAMING_SNAKE_CASE__  )

						self.assertTrue(torch.allclose(outputs.logits[0, :3]							,			SCREAMING_SNAKE_CASE__							,			atol=1e-4  )  )






		@slow
		def  __A       (		self		:							List[str]  )     ->							Any:
						__lowerCamelCase    					=	MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''  )
						__lowerCamelCase    					=	model.to(SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''  )

						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(images=SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )

						# forward pass
						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	outputs.logits

						# verify the logits
						__lowerCamelCase    					=	torch.Size((1, 21, 32, 32)  )
						self.assertEqual(logits.shape							,			SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	torch.tensor(
						    [
						        [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
						        [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
						        [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
						    ]							,			device=SCREAMING_SNAKE_CASE__							,			)

						self.assertTrue(torch.allclose(logits[0, :3, :3, :3]							,			SCREAMING_SNAKE_CASE__							,			atol=1e-4  )  )







		@slow
		def  __A       (		self		:							Union[str, Any]  )     ->							Optional[int]:
						__lowerCamelCase    					=	MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''  )
						__lowerCamelCase    					=	model.to(SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''  )

						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(images=SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )

						# forward pass
						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	outputs.logits.detach().cpu()

						__lowerCamelCase    					=	image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE__							,			target_sizes=[(50, 60)]  )
						__lowerCamelCase    					=	torch.Size((50, 60)  )
						self.assertEqual(segmentation[0].shape							,			SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	image_processor.post_process_semantic_segmentation(outputs=SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	torch.Size((32, 32)  )
						self.assertEqual(segmentation[0].shape							,			SCREAMING_SNAKE_CASE__  )





import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
					from sagemaker import Session, TrainingJobAnalytics
					from sagemaker.huggingface import HuggingFace

@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""  ,	"""False"""			)			) is not True  ,	reason="""Skipping test because should only be run when releasing minor transformers version"""  ,	)
@pytest.mark.usefixtures("""sm_env"""			)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue_model_parallelism.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
    ]			)
class    lowerCAmelCase__     (	unittest.TestCase			):






		def  __A       (		self		:							Optional[int]  )     ->							List[Any]:
						if self.framework == "pytorch":
										subprocess.run(
										    f'''cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py'''.split()							,			encoding='''utf-8'''							,			check=SCREAMING_SNAKE_CASE__							,			)
						assert hasattr(self							,			'''env'''  )






		def  __A       (		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							Optional[Any]:
						# configuration for running training on smdistributed Model Parallel
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''processes_per_host''': 8,
						}
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''parameters''': {
						        '''microbatches''': 4,
						        '''placement_strategy''': '''spread''',
						        '''pipeline''': '''interleaved''',
						        '''optimize''': '''speed''',
						        '''partitions''': 4,
						        '''ddp''': True,
						    },
						}

						__lowerCamelCase    					=	{'''smdistributed''': {'''modelparallel''': smp_options}, '''mpi''': mpi_options}

						__lowerCamelCase    					=	'''trainer''' if self.script == '''run_glue.py''' else '''smtrainer'''
						# creates estimator
						return HuggingFace(
						    entry_point=self.script							,			source_dir=self.env.test_path							,			role=self.env.role							,			image_uri=self.env.image_uri							,			base_job_name=f'''{self.env.base_job_name}-{instance_count}-smp-{name_extension}'''							,			instance_count=SCREAMING_SNAKE_CASE__							,			instance_type=self.instance_type							,			debugger_hook_config=SCREAMING_SNAKE_CASE__							,			hyperparameters={
						        **self.env.hyperparameters,
						        '''model_name_or_path''': self.model_name_or_path,
						        '''max_steps''': 5_00,
						    }							,			metric_definitions=self.env.metric_definitions							,			distribution=SCREAMING_SNAKE_CASE__							,			py_version='''py36'''							,			)






		def  __A       (		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							List[str]:
						TrainingJobAnalytics(SCREAMING_SNAKE_CASE__  ).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv'''  )







		@parameterized.expand([(1,)]  )
		def  __A       (		self		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]  )     ->							Optional[int]:
						# create estimator
						__lowerCamelCase    					=	self.create_estimator(SCREAMING_SNAKE_CASE__  )

						# run training
						estimator.fit()

						# result dataframe
						__lowerCamelCase    					=	TrainingJobAnalytics(estimator.latest_training_job.name  ).dataframe()

						# extract kpis
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']  )
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']  )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__lowerCamelCase    					=	(
						    Session().describe_training_job(estimator.latest_training_job.name  ).get('''TrainingTimeInSeconds'''							,			99_99_99  )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy  )
						assert all(t <= self.results['''eval_loss'''] for t in eval_loss  )

						# dump tests result into json file to share in PR
						with open(f'''{estimator.latest_training_job.name}.json'''							,			'''w'''  ) as outfile:
										json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}							,			SCREAMING_SNAKE_CASE__  )



