

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

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.config_class(**self.inputs_dict    )
     self.parent.assertTrue(hasattr(UpperCAmelCase_							,    'width_multiplier'    )    )




class      UpperCamelCase :

    def __init__(			self :      int							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Tuple=1_3							,    snake_case__ :      Dict=6_4							,    snake_case__ :      Dict=2							,    snake_case__ :      int=3							,    snake_case__ :      Any="swish"							,    snake_case__ :      str=3							,    snake_case__ :      Union[str, Any]=3_2							,    snake_case__ :      Any=0.1							,    snake_case__ :      Optional[int]=0.02							,    snake_case__ :      List[Any]=True							,    snake_case__ :      Tuple=True							,    snake_case__ :      Tuple=1_0							,    snake_case__ :      Optional[Any]=None							,    snake_case__ :      Any=0.25							,    snake_case__ :      Tuple=0.0							,    snake_case__ :      Union[str, Any]=0.0							,    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

     SCREAMING_SNAKE_CASE								=			self.get_config()

     return config, pixel_values, labels, pixel_labels

    def UpperCamelCase			(			self :      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 UpperCamelCase			(			self :      List[Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      List[Any]							,    snake_case__ :      int							,    snake_case__ :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MobileViTVaModel(config=UpperCAmelCase_    )
     model.to(UpperCAmelCase_    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Tuple							,    snake_case__ :      Any							,    snake_case__ :      Any							,    snake_case__ :      int							,    snake_case__ :      Dict    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      str							,    snake_case__ :      str							,    snake_case__ :      int							,    snake_case__ :      str							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_labels
     SCREAMING_SNAKE_CASE								=			MobileViTVaForSemanticSegmentation(UpperCAmelCase_    )
     model.to(UpperCAmelCase_    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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,
         )							,    )
     SCREAMING_SNAKE_CASE								=			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 :      Any    ):

     """simple docstring"""






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




@require_torch
class      UpperCamelCase (     __UpperCAmelCase	,			__UpperCAmelCase	,			unittest.TestCase   ):
    __UpperCamelCase         =(
        (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
        if is_torch_available()
        else ()
    )

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

    __UpperCamelCase         =False
    __UpperCamelCase         =False
    __UpperCamelCase         =False
    __UpperCamelCase         =False

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MobileViTVaModelTester(self    )
     SCREAMING_SNAKE_CASE								=			MobileViTVaConfigTester(self							,    config_class=UpperCAmelCase_							,    has_text_modality=UpperCAmelCase_    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     self.config_tester.run_common_tests()

    @unittest.skip(reason='MobileViTV2 does not use inputs_embeds'    )
    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     pass

    @unittest.skip(reason='MobileViTV2 does not support input and output embeddings'    )
    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     pass

    @unittest.skip(reason='MobileViTV2 does not output attentions'    )
    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     pass

    @require_torch_multi_gpu
    @unittest.skip(reason='Got `CUDA error: misaligned address` for tests after this one being run.'    )
    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     pass

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

     """simple docstring"""






     pass

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()

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

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

    def UpperCamelCase			(			self :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCAmelCase_    )

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     def check_hidden_states_output(snake_case__ :      List[str]							,    snake_case__ :      List[Any]							,    snake_case__ :      Any    ):
      SCREAMING_SNAKE_CASE								=			model_class(UpperCAmelCase_    )
      model.to(UpperCAmelCase_    )
      model.eval()

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

      SCREAMING_SNAKE_CASE								=			outputs.hidden_states

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

     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      SCREAMING_SNAKE_CASE								=			True
      check_hidden_states_output(UpperCAmelCase_							,    UpperCAmelCase_							,    UpperCAmelCase_    )

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

      check_hidden_states_output(UpperCAmelCase_							,    UpperCAmelCase_							,    UpperCAmelCase_    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_    )




    @slow
    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






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


def        __lowerCAmelCase    (       )		->							Tuple:
 '''simple docstring'''


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




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

    @cached_property
    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






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

    @slow
    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MobileViTVaForImageClassification.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256'    ).to(
         UpperCAmelCase_    )

     SCREAMING_SNAKE_CASE								=			self.default_image_processor
     SCREAMING_SNAKE_CASE								=			prepare_img()
     SCREAMING_SNAKE_CASE								=			image_processor(images=UpperCAmelCase_							,    return_tensors='pt'    ).to(UpperCAmelCase_    )

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

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

     SCREAMING_SNAKE_CASE								=			torch.tensor([-1.63_36E00, -7.32_04E-02, -5.18_83E-01]    ).to(UpperCAmelCase_    )

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

    @slow
    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3'    )
     SCREAMING_SNAKE_CASE								=			model.to(UpperCAmelCase_    )

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

     SCREAMING_SNAKE_CASE								=			prepare_img()
     SCREAMING_SNAKE_CASE								=			image_processor(images=UpperCAmelCase_							,    return_tensors='pt'    ).to(UpperCAmelCase_    )

     # forward pass
     with torch.no_grad():
      SCREAMING_SNAKE_CASE								=			model(**UpperCAmelCase_    )
     SCREAMING_SNAKE_CASE								=			outputs.logits

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

     SCREAMING_SNAKE_CASE								=			torch.tensor(
         [
             [[7.0_863, 7.1_525, 6.8_201], [6.6_931, 6.8_770, 6.8_933], [6.2_978, 7.0_366, 6.9_636]],
             [[-3.7_134, -3.6_712, -3.6_675], [-3.5_825, -3.3_549, -3.4_777], [-3.3_435, -3.3_979, -3.2_857]],
             [[-2.9_329, -2.8_003, -2.7_369], [-3.0_564, -2.4_780, -2.0_207], [-2.6_889, -1.9_298, -1.7_640]],
         ]							,    device=UpperCAmelCase_							,    )

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




    @slow
    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3'    )
     SCREAMING_SNAKE_CASE								=			model.to(UpperCAmelCase_    )

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

     SCREAMING_SNAKE_CASE								=			prepare_img()
     SCREAMING_SNAKE_CASE								=			image_processor(images=UpperCAmelCase_							,    return_tensors='pt'    ).to(UpperCAmelCase_    )

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

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

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

     SCREAMING_SNAKE_CASE								=			image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_    )
     SCREAMING_SNAKE_CASE								=			torch.Size((3_2, 3_2)    )
     self.assertEqual(segmentation[0].shape							,    UpperCAmelCase_    )









def        __lowerCAmelCase    (       _UpperCamelCase		:  str		,							_UpperCamelCase		:  int	)		->							list[str]:
 '''simple docstring'''


 return [sentence[i : i + ngram_size] for i in range(len(_UpperCamelCase	) - ngram_size + 1	)]


if __name__ == "__main__":
      from doctest import testmod

      testmod()

