"""simple docstring"""






from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import numpy as np
       import tensorflow as tf

       from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
       from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import MobileViTImageProcessor


class 					UpperCAmelCase		(_UpperCAmelCase						):



    """simple docstring"""

    def _snake_case    (     self							):
       lowercase__:     Dict      =   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 					UpperCAmelCase		:



    """simple docstring"""

    def __init__(     self					,    _UpperCAmelCase					,    _UpperCAmelCase=13					,    _UpperCAmelCase=32					,    _UpperCAmelCase=2					,    _UpperCAmelCase=3					,    _UpperCAmelCase=640					,    _UpperCAmelCase=4					,    _UpperCAmelCase="silu"					,    _UpperCAmelCase=3					,    _UpperCAmelCase=32					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=0.02					,    _UpperCAmelCase=True					,    _UpperCAmelCase=True					,    _UpperCAmelCase=10					,    _UpperCAmelCase=None					,    ):
       lowercase__:     List[str]      =   parent
       lowercase__:     Optional[int]      =   batch_size
       lowercase__:     Any      =   image_size
       lowercase__:     str      =   patch_size
       lowercase__:     Dict      =   num_channels
       lowercase__:     List[str]      =   last_hidden_size
       lowercase__:     List[str]      =   num_attention_heads
       lowercase__:     List[Any]      =   hidden_act
       lowercase__:     List[str]      =   conv_kernel_size
       lowercase__:     Any      =   output_stride
       lowercase__:     Union[str, Any]      =   hidden_dropout_prob
       lowercase__:     List[str]      =   attention_probs_dropout_prob
       lowercase__:     Any      =   classifier_dropout_prob
       lowercase__:     Optional[int]      =   use_labels
       lowercase__:     int      =   is_training
       lowercase__:     List[Any]      =   num_labels
       lowercase__:     List[Any]      =   initializer_range
       lowercase__:     Optional[Any]      =   scope






    def _snake_case    (     self							):
       lowercase__:     Union[str, Any]      =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

       lowercase__:     Tuple      =   None
       lowercase__:     Tuple      =   None
       if self.use_labels:
          lowercase__:     Optional[int]      =   ids_tensor([self.batch_size]					,    self.num_labels							)
          lowercase__:     int      =   ids_tensor([self.batch_size, self.image_size, self.image_size]					,    self.num_labels							)

       lowercase__:     List[Any]      =   self.get_config()

       return config, pixel_values, labels, pixel_labels






    def _snake_case    (     self							):
       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 _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							):
       lowercase__:     Any      =   TFMobileViTModel(config=_UpperCAmelCase							)
       lowercase__:     Tuple      =   model(_UpperCAmelCase					,    training=_UpperCAmelCase							)
       lowercase__:     Tuple      =   self.image_size // self.output_stride
       self.parent.assertEqual(
           result.last_hidden_state.shape					,    (self.batch_size, self.last_hidden_size, expected_height, expected_width)							)






    def _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							):
       lowercase__:     Optional[int]      =   self.num_labels
       lowercase__:     Tuple      =   TFMobileViTForImageClassification(_UpperCAmelCase							)
       lowercase__:     Tuple      =   model(_UpperCAmelCase					,    labels=_UpperCAmelCase					,    training=_UpperCAmelCase							)
       self.parent.assertEqual(result.logits.shape					,    (self.batch_size, self.num_labels)							)






    def _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							):
       lowercase__:     int      =   self.num_labels
       lowercase__:     Tuple      =   TFMobileViTForSemanticSegmentation(_UpperCAmelCase							)
       lowercase__:     Any      =   self.image_size // self.output_stride

       lowercase__:     Union[str, Any]      =   model(_UpperCAmelCase					,    training=_UpperCAmelCase							)
       self.parent.assertEqual(
           result.logits.shape					,    (self.batch_size, self.num_labels, expected_height, expected_width)							)

       lowercase__:     List[str]      =   model(_UpperCAmelCase					,    labels=_UpperCAmelCase					,    training=_UpperCAmelCase							)
       self.parent.assertEqual(
           result.logits.shape					,    (self.batch_size, self.num_labels, expected_height, expected_width)							)






    def _snake_case    (     self							):
       lowercase__:     Union[str, Any]      =   self.prepare_config_and_inputs()
       lowercase__,				lowercase__,				lowercase__,				lowercase__:     Any      =   config_and_inputs
       lowercase__:     int      =   {'''pixel_values''': pixel_values}
       return config, inputs_dict



@require_tf
class 					UpperCAmelCase		(_UpperCAmelCase					,_UpperCAmelCase					,unittest.TestCase						):



    """simple docstring"""

    _UpperCAmelCase			:Any							=				(
        (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
        if is_tf_available()
        else ()
    )
    _UpperCAmelCase			:Optional[int]							=				(
        {"feature-extraction": TFMobileViTModel, "image-classification": TFMobileViTForImageClassification}
        if is_tf_available()
        else {}
    )

    _UpperCAmelCase			:int							=				False
    _UpperCAmelCase			:Any							=				False
    _UpperCAmelCase			:str							=				False
    _UpperCAmelCase			:Optional[int]							=				False
    _UpperCAmelCase			:Any							=				False






    def _snake_case    (     self							):
       lowercase__:     Dict      =   TFMobileViTModelTester(self							)
       lowercase__:     Optional[int]      =   TFMobileViTConfigTester(self					,    config_class=_UpperCAmelCase					,    has_text_modality=_UpperCAmelCase							)






    def _snake_case    (     self							):
       self.config_tester.run_common_tests()






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






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






    @unittest.skip(reason='''MobileViT does not output attentions'''							)
    def _snake_case    (     self							):
       pass






    def _snake_case    (     self							):
       lowercase__,				lowercase__:     Tuple      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          lowercase__:     Optional[Any]      =   model_class(_UpperCAmelCase							)
          lowercase__:     Any      =   inspect.signature(model.call							)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          lowercase__:     Tuple      =   [*signature.parameters.keys()]

          lowercase__:     Dict      =   ['''pixel_values''']
          self.assertListEqual(arg_names[:1]					,    _UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     Optional[int]      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase							)






    def _snake_case    (     self							):
       def check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							):
          lowercase__:     Any      =   model_class(_UpperCAmelCase							)

          lowercase__:     Any      =   model(**self._prepare_for_class(_UpperCAmelCase					,    _UpperCAmelCase							)							)

          lowercase__:     Optional[Any]      =   outputs.hidden_states

          lowercase__:     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.
          lowercase__:     Union[str, Any]      =   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							)

       lowercase__,				lowercase__:     Any      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          lowercase__:     str      =   True
          check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							)

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

          check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     Optional[Any]      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase							)






    def _snake_case    (     self							):
       lowercase__:     List[str]      =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCAmelCase							)






    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices('''GPU'''							)							) == 0					,    reason='''TF does not support backprop for grouped convolutions on CPU.'''					,    )
    def _snake_case    (     self							):
       super().test_dataset_conversion()






    def _snake_case    (     self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase=2e-1					,    _UpperCAmelCase=2e-1							):
       self.assertTrue(np.allclose(_UpperCAmelCase					,    _UpperCAmelCase					,    atol=_UpperCAmelCase					,    rtol=_UpperCAmelCase							)							)






    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices('''GPU'''							)							) == 0					,    reason='''TF does not support backprop for grouped convolutions on CPU.'''					,    )
    @slow
    def _snake_case    (     self							):
       lowercase__,				lowercase__:     Tuple      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          # Since `TFMobileViTModel` cannot operate with the default `fit()` method.
          if model_class.__name__ != "TFMobileViTModel":
             lowercase__:     Dict      =   model_class(_UpperCAmelCase							)
             if getattr(_UpperCAmelCase					,    '''hf_compute_loss'''					,    _UpperCAmelCase							):
                super().test_keras_fit()






    def _snake_case    (     self							):
       lowercase__,				lowercase__:     List[Any]      =   self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
          # set an ignore index to correctly test the masked loss used in
          # `TFMobileViTForSemanticSegmentation`.
          if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
             lowercase__:     Tuple      =   5

          lowercase__:     str      =   model_class(_UpperCAmelCase							)
          if getattr(_UpperCAmelCase					,    '''hf_compute_loss'''					,    _UpperCAmelCase							):
             # The number of elements in the loss should be the same as the number of elements in the label
             lowercase__:     List[str]      =   self._prepare_for_class(inputs_dict.copy()					,    _UpperCAmelCase					,    return_labels=_UpperCAmelCase							)
             lowercase__:     List[Any]      =   prepared_for_class[
                 sorted(prepared_for_class.keys() - inputs_dict.keys()					,    reverse=_UpperCAmelCase							)[0]
             ]
             lowercase__:     int      =   added_label.shape.as_list()[:1]

             # Test that model correctly compute the loss with kwargs
             lowercase__:     List[str]      =   self._prepare_for_class(inputs_dict.copy()					,    _UpperCAmelCase					,    return_labels=_UpperCAmelCase							)
             lowercase__:     str      =   {'''input_ids''', '''pixel_values''', '''input_features'''}
             lowercase__:     Tuple      =   possible_input_names.intersection(set(_UpperCAmelCase							)							).pop()
             lowercase__:     int      =   prepared_for_class.pop(_UpperCAmelCase							)

             lowercase__:     Any      =   model(_UpperCAmelCase					,    **_UpperCAmelCase							)[0]
             self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]							)

             # Test that model correctly compute the loss when we mask some positions
             lowercase__:     List[str]      =   self._prepare_for_class(inputs_dict.copy()					,    _UpperCAmelCase					,    return_labels=_UpperCAmelCase							)
             lowercase__:     Union[str, Any]      =   {'''input_ids''', '''pixel_values''', '''input_features'''}
             lowercase__:     Optional[int]      =   possible_input_names.intersection(set(_UpperCAmelCase							)							).pop()
             lowercase__:     List[Any]      =   prepared_for_class.pop(_UpperCAmelCase							)
             if "labels" in prepared_for_class:
                lowercase__:     int      =   prepared_for_class['''labels'''].numpy()
                if len(labels.shape							) > 1 and labels.shape[1] != 1:
                   # labels[0] = -100
                   lowercase__:     Optional[Any]      =   tf.convert_to_tensor(_UpperCAmelCase							)
                   lowercase__:     List[str]      =   model(_UpperCAmelCase					,    **_UpperCAmelCase							)[0]
                   self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]							)
                   self.assertTrue(not np.any(np.isnan(loss.numpy()							)							)							)

                # Test that model correctly compute the loss with a dict
             lowercase__:     Optional[Any]      =   self._prepare_for_class(inputs_dict.copy()					,    _UpperCAmelCase					,    return_labels=_UpperCAmelCase							)
             lowercase__:     Optional[int]      =   model(_UpperCAmelCase							)[0]
             self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]							)

             # Test that model correctly compute the loss with a tuple
             lowercase__:     Optional[int]      =   self._prepare_for_class(inputs_dict.copy()					,    _UpperCAmelCase					,    return_labels=_UpperCAmelCase							)

             # Get keys that were added with the _prepare_for_class function
             lowercase__:     Optional[int]      =   prepared_for_class.keys() - inputs_dict.keys()
             lowercase__:     List[str]      =   inspect.signature(model.call							).parameters
             lowercase__:     Union[str, Any]      =   list(signature.keys()							)

             # Create a dictionary holding the location of the tensors in the tuple
             lowercase__:     Tuple      =   {0: input_name}
             for label_key in label_keys:
                lowercase__:     Union[str, Any]      =   signature_names.index(_UpperCAmelCase							)
                lowercase__:     Dict      =   label_key
             lowercase__:     Dict      =   sorted(tuple_index_mapping.items()							)
             # Initialize a list with their default values, update the values and convert to a tuple
             lowercase__:     Optional[Any]      =   []

             for name in signature_names:
                if name != "kwargs":
                   list_input.append(signature[name].default							)

             for index, value in sorted_tuple_index_mapping:
                lowercase__:     int      =   prepared_for_class[value]

             lowercase__:     int      =   tuple(_UpperCAmelCase							)

             # Send to model
             lowercase__:     int      =   model(tuple_input[:-1]							)[0]

             self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]							)






    @slow
    def _snake_case    (     self							):
       for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase__:     Optional[int]      =   TFMobileViTModel.from_pretrained(_UpperCAmelCase							)
          self.assertIsNotNone(_UpperCAmelCase							)



def SCREAMING_SNAKE_CASE__    (     )						->						Dict:
   lowercase__:     Dict      =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
   return image




@require_tf
class 					UpperCAmelCase		(unittest.TestCase						):



    """simple docstring"""

    @slow
    def _snake_case    (     self							):
       lowercase__:     Any      =   TFMobileViTForImageClassification.from_pretrained('''apple/mobilevit-xx-small'''							)

       lowercase__:     str      =   MobileViTImageProcessor.from_pretrained('''apple/mobilevit-xx-small'''							)
       lowercase__:     Union[str, Any]      =   prepare_img()
       lowercase__:     Optional[Any]      =   image_processor(images=_UpperCAmelCase					,    return_tensors='''tf'''							)

       # forward pass
       lowercase__:     Dict      =   model(**_UpperCAmelCase					,    training=_UpperCAmelCase							)

       # verify the logits
       lowercase__:     List[str]      =   tf.TensorShape((1, 1000)							)
       self.assertEqual(outputs.logits.shape					,    _UpperCAmelCase							)

       lowercase__:     List[str]      =   tf.constant([-1.9_364, -1.2_327, -0.4_653]							)

       tf.debugging.assert_near(outputs.logits[0, :3]					,    _UpperCAmelCase					,    atol=1e-4					,    rtol=1e-0_4							)






    @slow
    def _snake_case    (     self							):
       # `from_pt` will be removed
       lowercase__:     Tuple      =   TFMobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''							)

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

       lowercase__:     Union[str, Any]      =   prepare_img()
       lowercase__:     List[str]      =   image_processor(images=_UpperCAmelCase					,    return_tensors='''tf'''							)

       # forward pass
       lowercase__:     Dict      =   model(inputs.pixel_values					,    training=_UpperCAmelCase							)
       lowercase__:     Union[str, Any]      =   outputs.logits

       # verify the logits
       lowercase__:     Union[str, Any]      =   tf.TensorShape((1, 21, 32, 32)							)
       self.assertEqual(logits.shape					,    _UpperCAmelCase							)

       lowercase__:     int      =   tf.constant(
           [
               [[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]],
           ]							)

       tf.debugging.assert_near(logits[0, :3, :3, :3]					,    _UpperCAmelCase					,    rtol=1e-4					,    atol=1e-4							)

"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A         =				logging.get_logger(__name__)

__A         =				{
    "uw-madison/yoso-4096": "https://huggingface.co/uw-madison/yoso-4096/resolve/main/config.json",
    # See all YOSO models at https://huggingface.co/models?filter=yoso
}


class 					UpperCAmelCase		(_UpperCAmelCase						):



    """simple docstring"""

    _UpperCAmelCase			:Tuple							=				"yoso"






    def __init__(     self					,    _UpperCAmelCase=50265					,    _UpperCAmelCase=768					,    _UpperCAmelCase=12					,    _UpperCAmelCase=12					,    _UpperCAmelCase=3072					,    _UpperCAmelCase="gelu"					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase=4096					,    _UpperCAmelCase=1					,    _UpperCAmelCase=0.02					,    _UpperCAmelCase=1e-1_2					,    _UpperCAmelCase="absolute"					,    _UpperCAmelCase=True					,    _UpperCAmelCase=9					,    _UpperCAmelCase=64					,    _UpperCAmelCase=None					,    _UpperCAmelCase=True					,    _UpperCAmelCase=True					,    _UpperCAmelCase=1					,    _UpperCAmelCase=0					,    _UpperCAmelCase=2					,    **_UpperCAmelCase					,    ):
       super().__init__(pad_token_id=_UpperCAmelCase					,    bos_token_id=_UpperCAmelCase					,    eos_token_id=_UpperCAmelCase					,    **_UpperCAmelCase							)

       lowercase__:     int      =   vocab_size
       lowercase__:     List[str]      =   max_position_embeddings
       lowercase__:     Tuple      =   hidden_size
       lowercase__:     Optional[Any]      =   num_hidden_layers
       lowercase__:     Optional[int]      =   num_attention_heads
       lowercase__:     List[Any]      =   intermediate_size
       lowercase__:     Union[str, Any]      =   hidden_act
       lowercase__:     int      =   hidden_dropout_prob
       lowercase__:     Optional[Any]      =   attention_probs_dropout_prob
       lowercase__:     List[Any]      =   initializer_range
       lowercase__:     Optional[Any]      =   type_vocab_size
       lowercase__:     int      =   layer_norm_eps
       lowercase__:     Optional[Any]      =   position_embedding_type
       lowercase__:     int      =   use_expectation
       lowercase__:     List[str]      =   hash_code_len
       lowercase__:     Union[str, Any]      =   num_hash
       lowercase__:     Any      =   conv_window
       lowercase__:     Dict      =   use_fast_hash
       lowercase__:     Dict      =   lsh_backward
