


import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, 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 MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
		from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import MobileViTImageProcessor



class   lowerCamelCase				(_snake_case      ):







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Dict:
							UpperCAmelCase_   :				str						=						self.config_class(**self.inputs_dict       )
							self.parent.assertTrue(hasattr(_UpperCamelCase    ,		'hidden_sizes'       )       )
							self.parent.assertTrue(hasattr(_UpperCamelCase    ,		'neck_hidden_sizes'       )       )
							self.parent.assertTrue(hasattr(_UpperCamelCase    ,		'num_attention_heads'       )       )





class   lowerCamelCase				:







		'''simple docstring'''



		def __init__(							self    ,		_UpperCamelCase    ,		_UpperCamelCase=1_3    ,		_UpperCamelCase=3_2    ,		_UpperCamelCase=2    ,		_UpperCamelCase=3    ,		_UpperCamelCase=6_4_0    ,		_UpperCamelCase=4    ,		_UpperCamelCase="silu"    ,		_UpperCamelCase=3    ,		_UpperCamelCase=3_2    ,		_UpperCamelCase=0.1    ,		_UpperCamelCase=0.1    ,		_UpperCamelCase=0.1    ,		_UpperCamelCase=0.02    ,		_UpperCamelCase=True    ,		_UpperCamelCase=True    ,		_UpperCamelCase=1_0    ,		_UpperCamelCase=None    ,		)     ->							Any:
							UpperCAmelCase_   :				str						=						parent
							UpperCAmelCase_   :				List[Any]						=						batch_size
							UpperCAmelCase_   :				Optional[int]						=						image_size
							UpperCAmelCase_   :				Optional[Any]						=						patch_size
							UpperCAmelCase_   :				List[str]						=						num_channels
							UpperCAmelCase_   :				Union[str, Any]						=						last_hidden_size
							UpperCAmelCase_   :				Union[str, Any]						=						num_attention_heads
							UpperCAmelCase_   :				List[Any]						=						hidden_act
							UpperCAmelCase_   :				int						=						conv_kernel_size
							UpperCAmelCase_   :				Dict						=						output_stride
							UpperCAmelCase_   :				Optional[Any]						=						hidden_dropout_prob
							UpperCAmelCase_   :				List[Any]						=						attention_probs_dropout_prob
							UpperCAmelCase_   :				List[Any]						=						classifier_dropout_prob
							UpperCAmelCase_   :				Any						=						use_labels
							UpperCAmelCase_   :				Any						=						is_training
							UpperCAmelCase_   :				Any						=						num_labels
							UpperCAmelCase_   :				str						=						initializer_range
							UpperCAmelCase_   :				Tuple						=						scope



		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				Any						=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

							UpperCAmelCase_   :				str						=						None
							UpperCAmelCase_   :				str						=						None
							if self.use_labels:
												UpperCAmelCase_   :				Any						=						ids_tensor([self.batch_size]    ,		self.num_labels       )
												UpperCAmelCase_   :				Tuple						=						ids_tensor([self.batch_size, self.image_size, self.image_size]    ,		self.num_labels       )

							UpperCAmelCase_   :				Optional[int]						=						self.get_config()

							return config, pixel_values, labels, pixel_labels



		def  __UpperCAmelCase      (							self       )     ->							Optional[Any]:
							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    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )     ->							Dict:
							UpperCAmelCase_   :				List[str]						=						MobileViTModel(config=_UpperCamelCase       )
							model.to(_UpperCamelCase       )
							model.eval()
							UpperCAmelCase_   :				List[str]						=						model(_UpperCamelCase       )
							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  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )     ->							Union[str, Any]:
							UpperCAmelCase_   :				Union[str, Any]						=						self.num_labels
							UpperCAmelCase_   :				Dict						=						MobileViTForImageClassification(_UpperCamelCase       )
							model.to(_UpperCamelCase       )
							model.eval()
							UpperCAmelCase_   :				List[Any]						=						model(_UpperCamelCase    ,		labels=_UpperCamelCase       )
							self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_labels)       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )     ->							Tuple:
							UpperCAmelCase_   :				Tuple						=						self.num_labels
							UpperCAmelCase_   :				Optional[int]						=						MobileViTForSemanticSegmentation(_UpperCamelCase       )
							model.to(_UpperCamelCase       )
							model.eval()
							UpperCAmelCase_   :				List[Any]						=						model(_UpperCamelCase       )
							self.parent.assertEqual(
							    result.logits.shape    ,		(
							        self.batch_size,
							        self.num_labels,
							        self.image_size // self.output_stride,
							        self.image_size // self.output_stride,
							    )    ,		)
							UpperCAmelCase_   :				Union[str, Any]						=						model(_UpperCamelCase    ,		labels=_UpperCamelCase       )
							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  __UpperCAmelCase      (							self       )     ->							List[str]:
							UpperCAmelCase_   :				int						=						self.prepare_config_and_inputs()
							UpperCAmelCase_   :				Optional[Any]						=						config_and_inputs
							UpperCAmelCase_   :				Dict						=						{'pixel_values': pixel_values}
							return config, inputs_dict





@require_torch
class   lowerCamelCase				(_snake_case			,  _snake_case			,  unittest.TestCase      ):







		'''simple docstring'''



		_snake_case   :   Optional[Any]							=  (
		    (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
		    if is_torch_available()
		    else ()
		)
		_snake_case   :   int							=  (
		    {
		        '''feature-extraction''': MobileViTModel,
		        '''image-classification''': MobileViTForImageClassification,
		        '''image-segmentation''': MobileViTForSemanticSegmentation,
		    }
		    if is_torch_available()
		    else {}
		)

		_snake_case   :   Optional[int]							=  False
		_snake_case   :   Optional[Any]							=  False
		_snake_case   :   List[str]							=  False
		_snake_case   :   str							=  False



		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							UpperCAmelCase_   :				Any						=						MobileViTModelTester(self       )
							UpperCAmelCase_   :				str						=						MobileViTConfigTester(self    ,		config_class=_UpperCamelCase    ,		has_text_modality=_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							self.config_tester.run_common_tests()



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



		@unittest.skip(reason='MobileViT does not support input and output embeddings'       )
		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							pass



		@unittest.skip(reason='MobileViT does not output attentions'       )
		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							pass



		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				Tuple						=						self.model_tester.prepare_config_and_inputs_for_common()

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

												UpperCAmelCase_   :				Dict						=						['pixel_values']
												self.assertListEqual(arg_names[:1]    ,		_UpperCamelCase       )



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



		def  __UpperCAmelCase      (							self       )     ->							Any:
							UpperCAmelCase_   :				str						=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							List[str]:
							def check_hidden_states_output(_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       ):
												UpperCAmelCase_   :				Optional[Any]						=						model_class(_UpperCamelCase       )
												model.to(_UpperCamelCase       )
												model.eval()

												with torch.no_grad():
																	UpperCAmelCase_   :				int						=						model(**self._prepare_for_class(_UpperCamelCase    ,		_UpperCamelCase       )       )

												UpperCAmelCase_   :				List[str]						=						outputs.hidden_states

												UpperCAmelCase_   :				Optional[int]						=						5
												self.assertEqual(len(_UpperCamelCase       )    ,		_UpperCamelCase       )

												# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
												# with the width and height being successively divided by 2.
												UpperCAmelCase_   :				Optional[int]						=						2
												for i in range(len(_UpperCamelCase       )       ):
																	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_   :				List[Any]						=						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												UpperCAmelCase_   :				Optional[Any]						=						True
												check_hidden_states_output(_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )

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

												check_hidden_states_output(_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							Union[str, Any]:
							UpperCAmelCase_   :				Dict						=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							UpperCAmelCase_   :				Optional[int]						=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCamelCase       )



		@slow
		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												UpperCAmelCase_   :				Optional[Any]						=						MobileViTModel.from_pretrained(_UpperCamelCase       )
												self.assertIsNotNone(_UpperCamelCase       )





def 					lowercase__				(					):


					'''simple docstring'''


					UpperCAmelCase_   :				str						=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
					return image






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







		'''simple docstring'''



		@cached_property
		def  __UpperCAmelCase      (							self       )     ->							str:
							return MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'       ) if is_vision_available() else None



		@slow
		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							UpperCAmelCase_   :				List[Any]						=						MobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'       ).to(_UpperCamelCase       )

							UpperCAmelCase_   :				Dict						=						self.default_image_processor
							UpperCAmelCase_   :				List[str]						=						prepare_img()
							UpperCAmelCase_   :				Union[str, Any]						=						image_processor(images=_UpperCamelCase    ,		return_tensors='pt'       ).to(_UpperCamelCase       )

							# forward pass
							with torch.no_grad():
												UpperCAmelCase_   :				int						=						model(**_UpperCamelCase       )

							# verify the logits
							UpperCAmelCase_   :				Any						=						torch.Size((1, 1_0_0_0)       )
							self.assertEqual(outputs.logits.shape    ,		_UpperCamelCase       )

							UpperCAmelCase_   :				Optional[Any]						=						torch.tensor([-1.93_64, -1.23_27, -0.46_53]       ).to(_UpperCamelCase       )

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



		@slow
		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							UpperCAmelCase_   :				int						=						MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )
							UpperCAmelCase_   :				Tuple						=						model.to(_UpperCamelCase       )

							UpperCAmelCase_   :				Optional[int]						=						MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )

							UpperCAmelCase_   :				str						=						prepare_img()
							UpperCAmelCase_   :				List[str]						=						image_processor(images=_UpperCamelCase    ,		return_tensors='pt'       ).to(_UpperCamelCase       )

							# forward pass
							with torch.no_grad():
												UpperCAmelCase_   :				Optional[int]						=						model(**_UpperCamelCase       )
							UpperCAmelCase_   :				List[str]						=						outputs.logits

							# verify the logits
							UpperCAmelCase_   :				Any						=						torch.Size((1, 2_1, 3_2, 3_2)       )
							self.assertEqual(logits.shape    ,		_UpperCamelCase       )

							UpperCAmelCase_   :				List[Any]						=						torch.tensor(
							    [
							        [[6.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
							        [[-10.68_69, -10.32_50, -10.34_71], [-10.42_28, -9.98_68, -9.71_32], [-11.04_05, -11.02_21, -10.73_18]],
							        [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
							    ]    ,		device=_UpperCamelCase    ,		)

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



		@slow
		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				Any						=						MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )
							UpperCAmelCase_   :				List[Any]						=						model.to(_UpperCamelCase       )

							UpperCAmelCase_   :				List[Any]						=						MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'       )

							UpperCAmelCase_   :				Tuple						=						prepare_img()
							UpperCAmelCase_   :				List[str]						=						image_processor(images=_UpperCamelCase    ,		return_tensors='pt'       ).to(_UpperCamelCase       )

							# forward pass
							with torch.no_grad():
												UpperCAmelCase_   :				Dict						=						model(**_UpperCamelCase       )

							UpperCAmelCase_   :				List[str]						=						outputs.logits.detach().cpu()

							UpperCAmelCase_   :				Any						=						image_processor.post_process_semantic_segmentation(outputs=_UpperCamelCase    ,		target_sizes=[(5_0, 6_0)]       )
							UpperCAmelCase_   :				Any						=						torch.Size((5_0, 6_0)       )
							self.assertEqual(segmentation[0].shape    ,		_UpperCamelCase       )

							UpperCAmelCase_   :				int						=						image_processor.post_process_semantic_segmentation(outputs=_UpperCamelCase       )
							UpperCAmelCase_   :				List[str]						=						torch.Size((3_2, 3_2)       )
							self.assertEqual(segmentation[0].shape    ,		_UpperCamelCase       )














import argparse

import torch

from transformers import RemBertConfig, RemBertModel, load_tf_weights_in_rembert
from transformers.utils import logging


logging.set_verbosity_info()




def 					lowercase__				(					__snake_case   :							Optional[Any]					,	__snake_case   :							Any					,	__snake_case   :							int  ):


					'''simple docstring'''


					UpperCAmelCase_   :				Union[str, Any]						=						RemBertConfig.from_json_file(__snake_case  )
					print('Building PyTorch model from configuration: {}'.format(str(__snake_case  )  )  )
					UpperCAmelCase_   :				Dict						=						RemBertModel(__snake_case  )

					# Load weights from tf checkpoint
					load_tf_weights_in_rembert(__snake_case					,	__snake_case					,	__snake_case  )

					# Save pytorch-model
					print('Save PyTorch model to {}'.format(__snake_case  )  )
					torch.save(model.state_dict()					,	__snake_case  )


if __name__ == "__main__":
		__UpperCAmelCase     							=						argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
		)
		parser.add_argument(
		    '--rembert_config_file',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        'The config json file corresponding to the pre-trained RemBERT model. \n'
		        'This specifies the model architecture.'
		    ),
		)
		parser.add_argument(
		    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
		)
		__UpperCAmelCase     							=						parser.parse_args()
		convert_rembert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.rembert_config_file, args.pytorch_dump_path)






