

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   UpperCamelCase			(  SCREAMING_SNAKE_CASE ):

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.config_class(**self.inputs_dict     )
     self.parent.assertTrue(hasattr(snake_case__				,     'hidden_sizes'     )     )
     self.parent.assertTrue(hasattr(snake_case__				,     'neck_hidden_sizes'     )     )
     self.parent.assertTrue(hasattr(snake_case__				,     'num_attention_heads'     )     )




class   UpperCamelCase			:

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

     """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		    =      last_hidden_size
     SCREAMING_SNAKE_CASE		    =      num_attention_heads
     SCREAMING_SNAKE_CASE		    =      hidden_act
     SCREAMING_SNAKE_CASE		    =      conv_kernel_size
     SCREAMING_SNAKE_CASE		    =      output_stride
     SCREAMING_SNAKE_CASE		    =      hidden_dropout_prob
     SCREAMING_SNAKE_CASE		    =      attention_probs_dropout_prob
     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

    def UpperCamelCase      (      self			:   Optional[int]     ):

     """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			:   Optional[int]     ):

     """simple docstring"""






     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			:   str				,     snake_case__			:   Union[str, Any]				,     snake_case__			:   Optional[int]				,     snake_case__			:   Any				,     snake_case__			:   Optional[int]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      MobileViTModel(config=snake_case__     )
     model.to(snake_case__     )
     model.eval()
     SCREAMING_SNAKE_CASE		    =      model(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			:   List[Any]				,     snake_case__			:   Union[str, Any]				,     snake_case__			:   Union[str, Any]				,     snake_case__			:   Tuple				,     snake_case__			:   Any     ):

     """simple docstring"""






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

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.num_labels
     SCREAMING_SNAKE_CASE		    =      MobileViTForSemanticSegmentation(snake_case__     )
     model.to(snake_case__     )
     model.eval()
     SCREAMING_SNAKE_CASE		    =      model(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,
         )				,     )
     SCREAMING_SNAKE_CASE		    =      model(snake_case__				,     labels=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			:   List[str]     ):

     """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			(  SCREAMING_SNAKE_CASE						,      SCREAMING_SNAKE_CASE						,      unittest.TestCase ):
    __UpperCamelCase    =(
        (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
        if is_torch_available()
        else ()
    )
    __UpperCamelCase    =(
        {
            "feature-extraction": MobileViTModel,
            "image-classification": MobileViTForImageClassification,
            "image-segmentation": MobileViTForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

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

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      MobileViTModelTester(self     )
     SCREAMING_SNAKE_CASE		    =      MobileViTConfigTester(self				,     config_class=snake_case__				,     has_text_modality=snake_case__     )

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     self.config_tester.run_common_tests()

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

     """simple docstring"""






     pass

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

     """simple docstring"""






     pass

    @unittest.skip(reason='MobileViT does not output attentions'     )
    def UpperCamelCase      (      self			:   List[Any]     ):

     """simple docstring"""






     pass

    def UpperCamelCase      (      self			:   Tuple     ):

     """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(snake_case__     )
      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]				,     snake_case__     )

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

     """simple docstring"""






     pass

    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






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

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






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

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

      SCREAMING_SNAKE_CASE		    =      outputs.hidden_states

      SCREAMING_SNAKE_CASE		    =      5
      self.assertEqual(len(snake_case__     )				,     snake_case__     )

      # MobileViT'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(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     )

     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(snake_case__				,     snake_case__				,     snake_case__     )

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

      check_hidden_states_output(snake_case__				,     snake_case__				,     snake_case__     )

    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(*snake_case__     )

    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(*snake_case__     )




    @slow
    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      SCREAMING_SNAKE_CASE		    =      MobileViTModel.from_pretrained(snake_case__     )
      self.assertIsNotNone(snake_case__     )


def 							__lowerCAmelCase     (							)		->				Optional[Any]:
 '''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/mobilevit-xx-small'     ) if is_vision_available() else None

    @slow
    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      MobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'     ).to(snake_case__     )

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

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

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

     SCREAMING_SNAKE_CASE		    =      torch.tensor([-1.9_364, -1.2_327, -0.4_653]     ).to(snake_case__     )

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

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )
     SCREAMING_SNAKE_CASE		    =      model.to(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )

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

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

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

     SCREAMING_SNAKE_CASE		    =      torch.tensor(
         [
             [[6.9_713, 6.9_786, 7.2_422], [7.2_893, 7.2_825, 7.4_446], [7.6_580, 7.8_797, 7.9_420]],
             [[-10.6_869, -10.3_250, -10.3_471], [-10.4_228, -9.9_868, -9.7_132], [-11.0_405, -11.0_221, -10.7_318]],
             [[-3.3_089, -2.8_539, -2.6_740], [-3.2_706, -2.5_621, -2.5_108], [-3.2_534, -2.6_615, -2.6_651]],
         ]				,     device=snake_case__				,     )

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




    @slow
    def UpperCamelCase      (      self			:   List[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )
     SCREAMING_SNAKE_CASE		    =      model.to(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )

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

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

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

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

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




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


a_     :    str     		=   [
    "EAGER",
    "AOT_EAGER",
    "INDUCTOR",
    "NVFUSER",
    "AOT_NVFUSER",
    "AOT_CUDAGRAPHS",
    "OFI",
    "FX2TRT",
    "ONNXRT",
    "IPEX",
]


def 							__lowerCAmelCase     (							_UpperCamelCase		:	int		,				_UpperCamelCase		:	List[Any]=None		,				_UpperCamelCase		:	Dict=None		,				_UpperCamelCase		:	str=None						)		->				Any:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      True
 while ask_again:
  SCREAMING_SNAKE_CASE		    =      input(_UpperCamelCase						)
  try:
   if default is not None and len(_UpperCamelCase						) == 0:
    return default
   return convert_value(_UpperCamelCase						) if convert_value is not None else result
  except Exception:
   if error_message is not None:
    print(_UpperCamelCase						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Optional[Any]		,				_UpperCamelCase		:	Union[str, Any]=[]		,				_UpperCamelCase		:	Optional[int]=None		,				_UpperCamelCase		:	Optional[Any]=0						)		->				List[Any]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      BulletMenu(_UpperCamelCase		,				_UpperCamelCase						)
 SCREAMING_SNAKE_CASE		    =      menu.run(default_choice=_UpperCamelCase						)
 return convert_value(_UpperCamelCase						) if convert_value is not None else result


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


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return ComputeEnvironment(['LOCAL_MACHINE', 'AMAZON_SAGEMAKER'][value]						)


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


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return DistributedType(['NO', 'MULTI_CPU', 'MULTI_XPU', 'MULTI_GPU', 'MULTI_NPU', 'TPU'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	List[Any]						)		->				List[str]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return DynamoBackend(DYNAMO_BACKENDS[value]						).value


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Union[str, Any]						)		->				Dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return PrecisionType(['no', 'fp16', 'bf16', 'fp8'][value]						)


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


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return SageMakerDistributedType(['NO', 'DATA_PARALLEL', 'MODEL_PARALLEL'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Dict						)		->				str:
 '''simple docstring'''


 return {"yes": True, "no": False}[value.lower()]




class   UpperCamelCase			(  argparse.RawDescriptionHelpFormatter ):

    def UpperCamelCase      (      self			:   Union[str, Any]				,     snake_case__			:   Tuple				,     snake_case__			:   Optional[Any]				,     snake_case__			:   List[str]				,     snake_case__			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      super()._format_usage(snake_case__				,     snake_case__				,     snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      usage.replace('<command> [<args>] '				,     ''     )
     return usage

