


"""simple docstring"""






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 ):

     '''simple docstring'''


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







class 			__lowerCamelCase     :

     '''simple docstring'''


     def __init__(					self			, __UpperCAmelCase			, __UpperCAmelCase=13			, __UpperCAmelCase=64			, __UpperCAmelCase=2			, __UpperCAmelCase=3			, __UpperCAmelCase="swish"			, __UpperCAmelCase=3			, __UpperCAmelCase=32			, __UpperCAmelCase=0.1			, __UpperCAmelCase=0.02			, __UpperCAmelCase=True			, __UpperCAmelCase=True			, __UpperCAmelCase=10			, __UpperCAmelCase=None			, __UpperCAmelCase=0.25			, __UpperCAmelCase=0.0			, __UpperCAmelCase=0.0			, )				->							Optional[Any]:
          _a					= parent
          _a					= batch_size
          _a					= image_size
          _a					= patch_size
          _a					= num_channels
          _a					= make_divisible(512 * width_multiplier			, divisor=8							)
          _a					= hidden_act
          _a					= conv_kernel_size
          _a					= output_stride
          _a					= classifier_dropout_prob
          _a					= use_labels
          _a					= is_training
          _a					= num_labels
          _a					= initializer_range
          _a					= scope
          _a					= width_multiplier
          _a					= ffn_dropout
          _a					= attn_dropout


     def 							_UpperCAmelCase					(					self							)				->							Optional[Any]:
          _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							)				->							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 							_UpperCAmelCase					(					self			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase							)				->							Any:
          _a					= MobileViTVaModel(config=SCREAMING_SNAKE_CASE__							)
          model.to(SCREAMING_SNAKE_CASE__							)
          model.eval()
          _a					= 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 							_UpperCAmelCase					(					self			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase							)				->							List[Any]:
          _a					= self.num_labels
          _a					= MobileViTVaForImageClassification(SCREAMING_SNAKE_CASE__							)
          model.to(SCREAMING_SNAKE_CASE__							)
          model.eval()
          _a					= model(SCREAMING_SNAKE_CASE__			, labels=SCREAMING_SNAKE_CASE__							)
          self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.num_labels)							)


     def 							_UpperCAmelCase					(					self			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase							)				->							Any:
          _a					= self.num_labels
          _a					= MobileViTVaForSemanticSegmentation(SCREAMING_SNAKE_CASE__							)
          model.to(SCREAMING_SNAKE_CASE__							)
          model.eval()
          _a					= 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,
              )			, )
          _a					= 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 							_UpperCAmelCase					(					self							)				->							Dict:
          _a					= self.prepare_config_and_inputs()
          _a		,					_a		,					_a		,					_a					= config_and_inputs
          _a					= {'''pixel_values''': pixel_values}
          return config, inputs_dict







@require_torch
class 			__lowerCamelCase     (    __lowercase     ,    __lowercase     ,    unittest.TestCase ):

     '''simple docstring'''
     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 							_UpperCAmelCase					(					self							)				->							Optional[Any]:
          _a					= MobileViTVaModelTester(self							)
          _a					= MobileViTVaConfigTester(self			, config_class=SCREAMING_SNAKE_CASE__			, has_text_modality=SCREAMING_SNAKE_CASE__							)


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


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


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


     @unittest.skip(reason='''MobileViTV2 does not output attentions'''							)
     def 							_UpperCAmelCase					(					self							)				->							Tuple:
          pass


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


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


     def 							_UpperCAmelCase					(					self							)				->							Any:
          _a		,					_a					= self.model_tester.prepare_config_and_inputs_for_common()

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

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


     def 							_UpperCAmelCase					(					self							)				->							Any:
          _a					= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__							)


     def 							_UpperCAmelCase					(					self							)				->							Optional[Any]:
          def check_hidden_states_output(__UpperCAmelCase			, __UpperCAmelCase			, __UpperCAmelCase							):
               _a					= model_class(SCREAMING_SNAKE_CASE__							)
               model.to(SCREAMING_SNAKE_CASE__							)
               model.eval()

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

               _a					= outputs.hidden_states

               _a					= 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.
               _a					= 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							)

          _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(SCREAMING_SNAKE_CASE__			, SCREAMING_SNAKE_CASE__			, SCREAMING_SNAKE_CASE__							)

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

               check_hidden_states_output(SCREAMING_SNAKE_CASE__			, SCREAMING_SNAKE_CASE__			, SCREAMING_SNAKE_CASE__							)


     def 							_UpperCAmelCase					(					self							)				->							Optional[Any]:
          _a					= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE__							)


     def 							_UpperCAmelCase					(					self							)				->							Tuple:
          _a					= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_SNAKE_CASE__							)





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







def 					A_      (      ):






     """simple docstring"""

     _a					= 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							)				->							Dict:
          return (
              MobileViTImageProcessor.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''							)
              if is_vision_available()
              else None
          )


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

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

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

          # verify the logits
          _a					= torch.Size((1, 1000)							)
          self.assertEqual(outputs.logits.shape			, SCREAMING_SNAKE_CASE__							)

          _a					= torch.tensor([-1.6_3_3_6e0_0, -7.3_2_0_4e-0_2, -5.1_8_8_3e-0_1]							).to(SCREAMING_SNAKE_CASE__							)

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


     @slow
     def 							_UpperCAmelCase					(					self							)				->							Any:
          _a					= MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''							)
          _a					= model.to(SCREAMING_SNAKE_CASE__							)

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

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

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

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

          _a					= 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 							_UpperCAmelCase					(					self							)				->							Optional[int]:
          _a					= MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''							)
          _a					= model.to(SCREAMING_SNAKE_CASE__							)

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

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

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

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

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

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

from typing import Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class    lowerCAmelCase__     (	__lowercase			):






		def __init__(		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							NestedDataStructureLike[PathLike]							,			SCREAMING_SNAKE_CASE__		:							Optional[NamedSplit] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[Features] = None							,			SCREAMING_SNAKE_CASE__		:							str = None							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							Optional[int] = None							,			**SCREAMING_SNAKE_CASE__		:							str							,			)     ->							Union[str, Any]:
						super().__init__(
						    SCREAMING_SNAKE_CASE__							,			split=SCREAMING_SNAKE_CASE__							,			features=SCREAMING_SNAKE_CASE__							,			cache_dir=SCREAMING_SNAKE_CASE__							,			keep_in_memory=SCREAMING_SNAKE_CASE__							,			streaming=SCREAMING_SNAKE_CASE__							,			num_proc=SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			)
						__lowerCamelCase    					=	path_or_paths if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ) else {self.split: path_or_paths}
						__lowerCamelCase    					=	Text(
						    cache_dir=SCREAMING_SNAKE_CASE__							,			data_files=SCREAMING_SNAKE_CASE__							,			features=SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			)







		def  __A       (		self		:							int  )     ->							Dict:
						# Build iterable dataset
						if self.streaming:
										__lowerCamelCase    					=	self.builder.as_streaming_dataset(split=self.split  )
						# Build regular (map-style) dataset
						else:
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None

										self.builder.download_and_prepare(
										    download_config=SCREAMING_SNAKE_CASE__							,			download_mode=SCREAMING_SNAKE_CASE__							,			verification_mode=SCREAMING_SNAKE_CASE__							,			base_path=SCREAMING_SNAKE_CASE__							,			num_proc=self.num_proc							,			)
										__lowerCamelCase    					=	self.builder.as_dataset(
										    split=self.split							,			verification_mode=SCREAMING_SNAKE_CASE__							,			in_memory=self.keep_in_memory  )
						return dataset



