





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 							_UpperCamelCase					(	SCREAMING_SNAKE_CASE_     ):

       def lowercase (		self: Union[str, Any]		)       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  self.config_class(**self.inputs_dict		)
              self.parent.assertTrue(hasattr(UpperCamelCase__					,   "width_multiplier"		)		)



class 							_UpperCamelCase					:

       def __init__(		self: List[str]					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: int=13					,   _SCREAMING_SNAKE_CASE: List[str]=64					,   _SCREAMING_SNAKE_CASE: Optional[int]=2					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=3					,   _SCREAMING_SNAKE_CASE: List[Any]="swish"					,   _SCREAMING_SNAKE_CASE: Optional[int]=3					,   _SCREAMING_SNAKE_CASE: Union[str, Any]=32					,   _SCREAMING_SNAKE_CASE: Dict=0.1					,   _SCREAMING_SNAKE_CASE: Tuple=0.02					,   _SCREAMING_SNAKE_CASE: Tuple=True					,   _SCREAMING_SNAKE_CASE: Tuple=True					,   _SCREAMING_SNAKE_CASE: Optional[int]=10					,   _SCREAMING_SNAKE_CASE: List[str]=None					,   _SCREAMING_SNAKE_CASE: str=0.25					,   _SCREAMING_SNAKE_CASE: int=0.0					,   _SCREAMING_SNAKE_CASE: Dict=0.0					,   )       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  parent
              UpperCamelCase_          =  batch_size
              UpperCamelCase_          =  image_size
              UpperCamelCase_          =  patch_size
              UpperCamelCase_          =  num_channels
              UpperCamelCase_          =  make_divisible(512 * width_multiplier					,   divisor=8		)
              UpperCamelCase_          =  hidden_act
              UpperCamelCase_          =  conv_kernel_size
              UpperCamelCase_          =  output_stride
              UpperCamelCase_          =  classifier_dropout_prob
              UpperCamelCase_          =  use_labels
              UpperCamelCase_          =  is_training
              UpperCamelCase_          =  num_labels
              UpperCamelCase_          =  initializer_range
              UpperCamelCase_          =  scope
              UpperCamelCase_          =  width_multiplier
              UpperCamelCase_          =  ffn_dropout
              UpperCamelCase_          =  attn_dropout

       def lowercase (		self: List[Any]		)       -> int:




              """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 lowercase (		self: Tuple		)       -> List[str]:




              """simple docstring"""


              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 lowercase (		self: Optional[Any]					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Any		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  MobileViTVaModel(config=UpperCamelCase__		)
              model.to(UpperCamelCase__		)
              model.eval()
              UpperCamelCase_          =  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 lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: Dict					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Union[str, Any]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  self.num_labels
              UpperCamelCase_          =  MobileViTVaForImageClassification(UpperCamelCase__		)
              model.to(UpperCamelCase__		)
              model.eval()
              UpperCamelCase_          =  model(UpperCamelCase__					,   labels=UpperCamelCase__		)
              self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.num_labels)		)

       def lowercase (		self: Tuple					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: Optional[int]		)       -> int:




              """simple docstring"""


              UpperCamelCase_          =  self.num_labels
              UpperCamelCase_          =  MobileViTVaForSemanticSegmentation(UpperCamelCase__		)
              model.to(UpperCamelCase__		)
              model.eval()
              UpperCamelCase_          =  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_          =  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 lowercase (		self: Optional[Any]		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  self.prepare_config_and_inputs()
              UpperCamelCase_          =  config_and_inputs
              UpperCamelCase_          =  {'''pixel_values''': pixel_values}
              return config, inputs_dict



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

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

       _UpperCamelCase				:					Dict         =   False
       _UpperCamelCase				:					Any         =   False
       _UpperCamelCase				:					Union[str, Any]         =   False
       _UpperCamelCase				:					Optional[Any]         =   False

       def lowercase (		self: Optional[int]		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  MobileViTVaModelTester(self		)
              UpperCamelCase_          =  MobileViTVaConfigTester(self					,   config_class=UpperCamelCase__					,   has_text_modality=UpperCamelCase__		)

       def lowercase (		self: Any		)       -> Optional[int]:




              """simple docstring"""


              self.config_tester.run_common_tests()

       @unittest.skip(reason="MobileViTV2 does not use inputs_embeds"		)
       def lowercase (		self: Dict		)       -> Dict:




              """simple docstring"""


              pass

       @unittest.skip(reason="MobileViTV2 does not support input and output embeddings"		)
       def lowercase (		self: Dict		)       -> int:




              """simple docstring"""


              pass

       @unittest.skip(reason="MobileViTV2 does not output attentions"		)
       def lowercase (		self: List[Any]		)       -> Optional[Any]:




              """simple docstring"""


              pass

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




              """simple docstring"""


              pass

       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."		)
       def lowercase (		self: List[Any]		)       -> str:




              """simple docstring"""


              pass

       def lowercase (		self: int		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_          =  self.model_tester.prepare_config_and_inputs_for_common()

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

                     UpperCamelCase_          =  ['''pixel_values''']
                     self.assertListEqual(arg_names[:1]					,   UpperCamelCase__		)

       def lowercase (		self: Optional[int]		)       -> Optional[int]:




              """simple docstring"""


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

       def lowercase (		self: Optional[Any]		)       -> Union[str, Any]:




              """simple docstring"""


              def check_hidden_states_output(_SCREAMING_SNAKE_CASE: List[str]					,   _SCREAMING_SNAKE_CASE: int					,   _SCREAMING_SNAKE_CASE: str		):
                     UpperCamelCase_          =  model_class(UpperCamelCase__		)
                     model.to(UpperCamelCase__		)
                     model.eval()

                     with torch.no_grad():
                            UpperCamelCase_          =  model(**self._prepare_for_class(UpperCamelCase__					,   UpperCamelCase__		)		)

                     UpperCamelCase_          =  outputs.hidden_states

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

                     # MobileViTV2'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(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_          =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     UpperCamelCase_          =  True
                     check_hidden_states_output(UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__		)

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

                     check_hidden_states_output(UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__		)

       def lowercase (		self: Any		)       -> List[Any]:




              """simple docstring"""


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

       def lowercase (		self: Tuple		)       -> Any:




              """simple docstring"""


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



       @slow
       def lowercase (		self: List[str]		)       -> Tuple:




              """simple docstring"""


              for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                     UpperCamelCase_          =  MobileViTVaModel.from_pretrained(UpperCamelCase__		)
                     self.assertIsNotNone(UpperCamelCase__		)




def  lowerCAmelCase_ (							)     ->		Dict:
       UpperCamelCase_          =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
       return image







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

       @cached_property
       def lowercase (		self: Optional[int]		)       -> Dict:




              """simple docstring"""


              return (
                  MobileViTImageProcessor.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"		)
                  if is_vision_available()
                  else None
              )

       @slow
       def lowercase (		self: Optional[int]		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  MobileViTVaForImageClassification.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"		).to(
                  UpperCamelCase__		)

              UpperCamelCase_          =  self.default_image_processor
              UpperCamelCase_          =  prepare_img()
              UpperCamelCase_          =  image_processor(images=UpperCamelCase__					,   return_tensors="pt"		).to(UpperCamelCase__		)

              # forward pass
              with torch.no_grad():
                     UpperCamelCase_          =  model(**UpperCamelCase__		)

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

              UpperCamelCase_          =  torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01]		).to(UpperCamelCase__		)

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

       @slow
       def lowercase (		self: Union[str, Any]		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"		)
              UpperCamelCase_          =  model.to(UpperCamelCase__		)

              UpperCamelCase_          =  MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"		)

              UpperCamelCase_          =  prepare_img()
              UpperCamelCase_          =  image_processor(images=UpperCamelCase__					,   return_tensors="pt"		).to(UpperCamelCase__		)

              # forward pass
              with torch.no_grad():
                     UpperCamelCase_          =  model(**UpperCamelCase__		)
              UpperCamelCase_          =  outputs.logits

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

              UpperCamelCase_          =  torch.tensor(
                  [
                      [[7.08_63, 7.15_25, 6.82_01], [6.69_31, 6.87_70, 6.89_33], [6.29_78, 7.03_66, 6.96_36]],
                      [[-3.71_34, -3.67_12, -3.66_75], [-3.58_25, -3.35_49, -3.47_77], [-3.34_35, -3.39_79, -3.28_57]],
                      [[-2.93_29, -2.80_03, -2.73_69], [-3.05_64, -2.47_80, -2.02_07], [-2.68_89, -1.92_98, -1.76_40]],
                  ]					,   device=UpperCamelCase__					,   )

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



       @slow
       def lowercase (		self: Optional[Any]		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"		)
              UpperCamelCase_          =  model.to(UpperCamelCase__		)

              UpperCamelCase_          =  MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"		)

              UpperCamelCase_          =  prepare_img()
              UpperCamelCase_          =  image_processor(images=UpperCamelCase__					,   return_tensors="pt"		).to(UpperCamelCase__		)

              # forward pass
              with torch.no_grad():
                     UpperCamelCase_          =  model(**UpperCamelCase__		)

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

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

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








from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class 							_UpperCamelCase					(	lowerCAmelCase_     ):

       def __init__(		self: List[Any]					,   _SCREAMING_SNAKE_CASE: pyspark.sql.DataFrame					,   _SCREAMING_SNAKE_CASE: Optional[NamedSplit] = None					,   _SCREAMING_SNAKE_CASE: Optional[Features] = None					,   _SCREAMING_SNAKE_CASE: bool = True					,   _SCREAMING_SNAKE_CASE: str = None					,   _SCREAMING_SNAKE_CASE: bool = False					,   _SCREAMING_SNAKE_CASE: str = None					,   _SCREAMING_SNAKE_CASE: bool = True					,   _SCREAMING_SNAKE_CASE: str = "arrow"					,   **_SCREAMING_SNAKE_CASE: Union[str, Any]					,   )       -> Any:




              """simple docstring"""


              super().__init__(
                  split=_SCREAMING_SNAKE_CASE					,   features=_SCREAMING_SNAKE_CASE					,   cache_dir=_SCREAMING_SNAKE_CASE					,   keep_in_memory=_SCREAMING_SNAKE_CASE					,   streaming=_SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE					,   )
              UpperCamelCase_          =  load_from_cache_file
              UpperCamelCase_          =  file_format
              UpperCamelCase_          =  Spark(
                  df=_SCREAMING_SNAKE_CASE					,   features=_SCREAMING_SNAKE_CASE					,   cache_dir=_SCREAMING_SNAKE_CASE					,   working_dir=_SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE					,   )



       def lowercase (		self: str		)       -> Optional[Any]:




              """simple docstring"""


              if self.streaming:
                     return self.builder.as_streaming_dataset(split=self.split		)
              UpperCamelCase_          =  None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
              self.builder.download_and_prepare(
                  download_mode=_SCREAMING_SNAKE_CASE					,   file_format=self._file_format					,   )
              return self.builder.as_dataset(split=self.split		)

