


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 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.config_class(**self.inputs_dict   )
       self.parent.assertTrue(hasattr(_a     ,		'''hidden_sizes'''   )   )
       self.parent.assertTrue(hasattr(_a     ,		'''neck_hidden_sizes'''   )   )
       self.parent.assertTrue(hasattr(_a     ,		'''num_attention_heads'''   )   )




class 		__magic_name__						:





     '''simple docstring'''



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

     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

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

       snake_case__						= self.get_config()

       return config, pixel_values, labels, pixel_labels

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       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 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:List[Any]     ,		_a:Dict     ,		_a:Optional[Any]     ,		_a:Union[str, Any]   ):
       snake_case__						= MobileViTModel(config=_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a   )
       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 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:List[Any]     ,		_a:Optional[Any]     ,		_a:int     ,		_a:int   ):
       snake_case__						= self.num_labels
       snake_case__						= MobileViTForImageClassification(_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a     ,		labels=_a   )
       self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.num_labels)   )

     def 	SCREAMING_SNAKE_CASE__						(					self:int     ,		_a:int     ,		_a:str     ,		_a:List[str]     ,		_a:Optional[int]   ):
       snake_case__						= self.num_labels
       snake_case__						= MobileViTForSemanticSegmentation(_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a   )
       self.parent.assertEqual(
           result.logits.shape     ,		(
               self.batch_size,
               self.num_labels,
               self.image_size // self.output_stride,
               self.image_size // self.output_stride,
           )     ,		)
       snake_case__						= model(_a     ,		labels=_a   )
       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 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__						= self.prepare_config_and_inputs()
       snake_case__						= config_and_inputs
       snake_case__						= {'pixel_values': pixel_values}
       return config, inputs_dict




@require_torch
class 		__magic_name__						(snake_case_   ,snake_case_   ,unittest.TestCase				):





     '''simple docstring'''



     __lowercase				:							Tuple	   =       (
         (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
         if is_torch_available()
         else ()
     )
     __lowercase				:							List[Any]	   =       (
         {
             '''feature-extraction''': MobileViTModel,
             '''image-classification''': MobileViTForImageClassification,
             '''image-segmentation''': MobileViTForSemanticSegmentation,
         }
         if is_torch_available()
         else {}
     )

     __lowercase				:							Optional[int]	   =       False
     __lowercase				:							Dict	   =       False
     __lowercase				:							List[Any]	   =       False
     __lowercase				:							Optional[int]	   =       False

     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__						= MobileViTModelTester(self   )
       snake_case__						= MobileViTConfigTester(self     ,		config_class=_a     ,		has_text_modality=_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]   ):
       self.config_tester.run_common_tests()

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

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

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

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

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

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

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

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):







       def check_hidden_states_output(_a:List[str]     ,		_a:List[str]     ,		_a:Optional[int]   ):
         snake_case__						= model_class(_a   )
         model.to(_a   )
         model.eval()

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

         snake_case__						= outputs.hidden_states

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

         # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
         # with the width and height being successively divided by 2.
         snake_case__						= 2
         for i in range(len(_a   )   ):
           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   )

       snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         snake_case__						= True
         check_hidden_states_output(_a     ,		_a     ,		_a   )

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

         check_hidden_states_output(_a     ,		_a     ,		_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_semantic_segmentation(*_a   )







     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[int]   ):







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







def SCREAMING_SNAKE_CASE	(			)    ->				Any:
  snake_case__						= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
  return image


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





     '''simple docstring'''



     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       return MobileViTImageProcessor.from_pretrained('''apple/mobilevit-xx-small'''   ) if is_vision_available() else None

     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= MobileViTForImageClassification.from_pretrained('''apple/mobilevit-xx-small'''   ).to(_a   )

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

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

       # verify the logits
       snake_case__						= torch.Size((1, 10_00)   )
       self.assertEqual(outputs.logits.shape     ,		_a   )

       snake_case__						= torch.tensor([-1.9364, -1.2327, -0.4653]   ).to(_a   )

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

     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''   )
       snake_case__						= model.to(_a   )

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

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

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

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

       snake_case__						= torch.tensor(
           [
               [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
               [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
               [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
           ]     ,		device=_a     ,		)

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







     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''   )
       snake_case__						= model.to(_a   )

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

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

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

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

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

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







'''simple docstring'''



from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


UpperCamelCase_    	=				logging.get_logger(__name__)


UpperCamelCase_    	=				TypeVar("DatasetType", Dataset, IterableDataset)




def   lowercase__(			__UpperCamelCase:					List[DatasetType]      ,__UpperCamelCase:					Optional[List[float]] = None      ,__UpperCamelCase:					Optional[int] = None      ,__UpperCamelCase:					Optional[DatasetInfo] = None      ,__UpperCamelCase:					Optional[NamedSplit] = None      ,__UpperCamelCase:					Literal["first_exhausted", "all_exhausted"] = "first_exhausted"      ,):

      """simple docstring"""






      from .arrow_dataset import Dataset
      from .iterable_dataset import IterableDataset

      if not datasets:
            raise ValueError('Unable to interleave an empty list of datasets.'    )
      for i, dataset in enumerate(__UpperCamelCase    ):
            if not isinstance(__UpperCamelCase      ,(Dataset, IterableDataset)    ):
                  if isinstance(__UpperCamelCase      ,(DatasetDict, IterableDatasetDict)    ):
                        if not dataset:
                              raise ValueError(
                                  f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
                                  'is an empty dataset dictionary.'    )
                        raise ValueError(
                            f"Dataset at position {i} has at least one split: {list(__UpperCamelCase    )}\n"
                            f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase    )    )}']"    )
                  raise ValueError(
                      f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase    ).__name__}."    )
            if i == 0:
                  SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						(
                      (Dataset, IterableDataset) if isinstance(__UpperCamelCase      ,__UpperCamelCase    ) else (IterableDataset, Dataset)
                  )
            elif not isinstance(__UpperCamelCase      ,__UpperCamelCase    ):
                  raise ValueError(
                      f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."    )
      if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
            raise ValueError(f"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."    )
      if dataset_type is Dataset:
            return _interleave_map_style_datasets(
                __UpperCamelCase      ,__UpperCamelCase      ,__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,stopping_strategy=__UpperCamelCase    )
      else:
            return _interleave_iterable_datasets(
                __UpperCamelCase      ,__UpperCamelCase      ,__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,stopping_strategy=__UpperCamelCase    )





def   lowercase__(			__UpperCamelCase:					List[DatasetType]      ,__UpperCamelCase:					Optional[DatasetInfo] = None      ,__UpperCamelCase:					Optional[NamedSplit] = None      ,__UpperCamelCase:					int = 0      ,):

      """simple docstring"""






      if not dsets:
            raise ValueError('Unable to concatenate an empty list of datasets.'    )
      for i, dataset in enumerate(__UpperCamelCase    ):
            if not isinstance(__UpperCamelCase      ,(Dataset, IterableDataset)    ):
                  if isinstance(__UpperCamelCase      ,(DatasetDict, IterableDatasetDict)    ):
                        if not dataset:
                              raise ValueError(
                                  f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
                                  'is an empty dataset dictionary.'    )
                        raise ValueError(
                            f"Dataset at position {i} has at least one split: {list(__UpperCamelCase    )}\n"
                            f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase    )    )}']"    )
                  raise ValueError(
                      f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase    ).__name__}."    )
            if i == 0:
                  SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      int    					=						(
                      (Dataset, IterableDataset) if isinstance(__UpperCamelCase      ,__UpperCamelCase    ) else (IterableDataset, Dataset)
                  )
            elif not isinstance(__UpperCamelCase      ,__UpperCamelCase    ):
                  raise ValueError(
                      f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."    )
      if dataset_type is Dataset:
            return _concatenate_map_style_datasets(__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,axis=__UpperCamelCase    )
      else:
            return _concatenate_iterable_datasets(__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,axis=__UpperCamelCase    )


