


import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 SCREAMING_SNAKE_CASE       :      str												= "platform"

 import jax
 import jax.numpy as jnp

 from transformers import BlenderbotTokenizer
 from transformers.models.blenderbot.modeling_flax_blenderbot import (
     FlaxBlenderbotForConditionalGeneration,
     FlaxBlenderbotModel,
     shift_tokens_right,
 )



def 				UpperCamelCase_(      lowerCamelCase_							, lowerCamelCase_							, lowerCamelCase_=None							, lowerCamelCase_=None							, lowerCamelCase_=None							, lowerCamelCase_=None							, lowerCamelCase_=None							, lowerCamelCase_=None							, )    ->						List[str]:
    if attention_mask is None:
        _lowercase       :							int      		=		np.where(input_ids != config.pad_token_id							, 1							, 0 )
    if decoder_attention_mask is None:
        _lowercase       :							Optional[Any]      		=		np.where(decoder_input_ids != config.pad_token_id							, 1							, 0 )
    if head_mask is None:
        _lowercase       :							Dict      		=		np.ones((config.encoder_layers, config.encoder_attention_heads) )
    if decoder_head_mask is None:
        _lowercase       :							Optional[int]      		=		np.ones((config.decoder_layers, config.decoder_attention_heads) )
    if cross_attn_head_mask is None:
        _lowercase       :							List[str]      		=		np.ones((config.decoder_layers, config.decoder_attention_heads) )
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": attention_mask,
    }






class _lowerCamelCase:
       def __init__(					self,      lowerCamelCase,      lowerCamelCase=13,      lowerCamelCase=7,      lowerCamelCase=True,      lowerCamelCase=False,      lowerCamelCase=99,      lowerCamelCase=16,      lowerCamelCase=2,      lowerCamelCase=4,      lowerCamelCase=4,      lowerCamelCase="gelu",      lowerCamelCase=0.1,      lowerCamelCase=0.1,      lowerCamelCase=32,      lowerCamelCase=2,      lowerCamelCase=1,      lowerCamelCase=0,      lowerCamelCase=0.0_2,      )       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		parent
           _lowercase       :							Any      		=		batch_size
           _lowercase       :							List[str]      		=		seq_length
           _lowercase       :							Union[str, Any]      		=		is_training
           _lowercase       :							Optional[Any]      		=		use_labels
           _lowercase       :							Tuple      		=		vocab_size
           _lowercase       :							Tuple      		=		hidden_size
           _lowercase       :							List[str]      		=		num_hidden_layers
           _lowercase       :							Union[str, Any]      		=		num_attention_heads
           _lowercase       :							Optional[int]      		=		intermediate_size
           _lowercase       :							List[Any]      		=		hidden_act
           _lowercase       :							Dict      		=		hidden_dropout_prob
           _lowercase       :							Any      		=		attention_probs_dropout_prob
           _lowercase       :							Union[str, Any]      		=		max_position_embeddings
           _lowercase       :							Tuple      		=		eos_token_id
           _lowercase       :							List[Any]      		=		pad_token_id
           _lowercase       :							Dict      		=		bos_token_id
           _lowercase       :							List[Any]      		=		initializer_range
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							str      		=		np.clip(ids_tensor([self.batch_size, self.seq_length - 1],      self.vocab_size),      3,      self.vocab_size)
           _lowercase       :							str      		=		np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1),      dtype=np.intaa)),      -1)

           _lowercase       :							Union[str, Any]      		=		shift_tokens_right(lowerCamelCase,      1,      2)

           _lowercase       :							Optional[Any]      		=		BlenderbotConfig(
               vocab_size=self.vocab_size,      d_model=self.hidden_size,      encoder_layers=self.num_hidden_layers,      decoder_layers=self.num_hidden_layers,      encoder_attention_heads=self.num_attention_heads,      decoder_attention_heads=self.num_attention_heads,      encoder_ffn_dim=self.intermediate_size,      decoder_ffn_dim=self.intermediate_size,      dropout=self.hidden_dropout_prob,      attention_dropout=self.attention_probs_dropout_prob,      max_position_embeddings=self.max_position_embeddings,      eos_token_id=self.eos_token_id,      bos_token_id=self.bos_token_id,      pad_token_id=self.pad_token_id,      initializer_range=self.initializer_range,      use_cache=lowerCamelCase,      )
           _lowercase       :							List[Any]      		=		prepare_blenderbot_inputs_dict(lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
           return config, inputs_dict
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Optional[Any]      		=		self.prepare_config_and_inputs()
           return config, inputs_dict
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							int      		=		20
           _lowercase       :							Any      		=		model_class_name(lowerCamelCase)

           _lowercase       :							Any      		=		model.encode(inputs_dict['input_ids'])

           _lowercase		,				_lowercase       :							Optional[int]      		=		(
               inputs_dict['decoder_input_ids'],
               inputs_dict['decoder_attention_mask'],
           )

           _lowercase       :							str      		=		model.init_cache(decoder_input_ids.shape[0],      lowerCamelCase,      lowerCamelCase)
           _lowercase       :							Tuple      		=		jnp.ones((decoder_input_ids.shape[0], max_decoder_length),      dtype='i4')

           _lowercase       :							Optional[int]      		=		jnp.broadcast_to(
               jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :],      (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1),      )
           _lowercase       :							Optional[Any]      		=		model.decode(
               decoder_input_ids[:, :-1],      lowerCamelCase,      decoder_attention_mask=lowerCamelCase,      past_key_values=lowerCamelCase,      decoder_position_ids=lowerCamelCase,      )

           _lowercase       :							Union[str, Any]      		=		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]],      dtype='i4')
           _lowercase       :							List[str]      		=		model.decode(
               decoder_input_ids[:, -1:],      lowerCamelCase,      decoder_attention_mask=lowerCamelCase,      past_key_values=outputs_cache.past_key_values,      decoder_position_ids=lowerCamelCase,      )

           _lowercase       :							Tuple      		=		model.decode(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							Union[str, Any]      		=		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
           self.parent.assertTrue(diff < 1E-3,      msg=F'''Max diff is {diff}''')




       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		20
           _lowercase       :							List[str]      		=		model_class_name(lowerCamelCase)

           _lowercase       :							int      		=		model.encode(inputs_dict['input_ids'])

           _lowercase		,				_lowercase       :							Any      		=		(
               inputs_dict['decoder_input_ids'],
               inputs_dict['decoder_attention_mask'],
           )

           _lowercase       :							str      		=		jnp.concatenate(
               [
                   decoder_attention_mask,
                   jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
               ],      axis=-1,      )

           _lowercase       :							Any      		=		model.init_cache(decoder_input_ids.shape[0],      lowerCamelCase,      lowerCamelCase)
           _lowercase       :							str      		=		jnp.broadcast_to(
               jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :],      (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1),      )

           _lowercase       :							Optional[Any]      		=		model.decode(
               decoder_input_ids[:, :-1],      lowerCamelCase,      decoder_attention_mask=lowerCamelCase,      past_key_values=lowerCamelCase,      decoder_position_ids=lowerCamelCase,      )
           _lowercase       :							str      		=		jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]],      dtype='i4')
           _lowercase       :							List[str]      		=		model.decode(
               decoder_input_ids[:, -1:],      lowerCamelCase,      past_key_values=outputs_cache.past_key_values,      decoder_attention_mask=lowerCamelCase,      decoder_position_ids=lowerCamelCase,      )

           _lowercase       :							List[Any]      		=		model.decode(lowerCamelCase,      lowerCamelCase,      decoder_attention_mask=lowerCamelCase)

           _lowercase       :							List[str]      		=		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
           self.parent.assertTrue(diff < 1E-3,      msg=F'''Max diff is {diff}''')






@require_flax
class _lowerCamelCase(						unittest.TestCase   ):
       lowercase_ :	str     =			99
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		np.array(
               [
                   [71, 82, 18, 33, 46, 91, 2],
                   [68, 34, 26, 58, 30, 82, 2],
                   [5, 97, 17, 39, 94, 40, 2],
                   [76, 83, 94, 25, 70, 78, 2],
                   [87, 59, 41, 35, 48, 66, 2],
                   [55, 13, 16, 58, 5, 2, 1],  # note padding
                   [64, 27, 31, 51, 12, 75, 2],
                   [52, 64, 86, 17, 83, 39, 2],
                   [48, 61, 9, 24, 71, 82, 2],
                   [26, 1, 60, 48, 22, 13, 2],
                   [21, 5, 62, 28, 14, 76, 2],
                   [45, 98, 37, 86, 59, 48, 2],
                   [70, 70, 50, 9, 28, 0, 2],
               ],      dtype=np.intaa,      )

           _lowercase       :							Optional[int]      		=		input_ids.shape[0]
           _lowercase       :							int      		=		BlenderbotConfig(
               vocab_size=self.vocab_size,      d_model=24,      encoder_layers=2,      decoder_layers=2,      encoder_attention_heads=2,      decoder_attention_heads=2,      encoder_ffn_dim=32,      decoder_ffn_dim=32,      max_position_embeddings=48,      eos_token_id=2,      pad_token_id=1,      bos_token_id=0,      )
           return config, input_ids, batch_size
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase		,				_lowercase		,				_lowercase       :							Dict      		=		self._get_config_and_data()
           _lowercase       :							Dict      		=		FlaxBlenderbotForConditionalGeneration(lowerCamelCase)
           _lowercase       :							Union[str, Any]      		=		lm_model(input_ids=lowerCamelCase)
           _lowercase       :							Tuple      		=		(batch_size, input_ids.shape[1], config.vocab_size)
           self.assertEqual(outputs['logits'].shape,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		BlenderbotConfig(
               vocab_size=self.vocab_size,      d_model=14,      encoder_layers=2,      decoder_layers=2,      encoder_attention_heads=2,      decoder_attention_heads=2,      encoder_ffn_dim=8,      decoder_ffn_dim=8,      max_position_embeddings=48,      )
           _lowercase       :							str      		=		FlaxBlenderbotForConditionalGeneration(lowerCamelCase)
           _lowercase       :							Optional[Any]      		=		np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]],      dtype=np.intaa)
           _lowercase       :							Optional[Any]      		=		np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]],      dtype=np.intaa)
           _lowercase       :							Dict      		=		lm_model(input_ids=lowerCamelCase,      decoder_input_ids=lowerCamelCase)
           _lowercase       :							Any      		=		(*summary.shape, config.vocab_size)
           self.assertEqual(outputs['logits'].shape,      lowerCamelCase)




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




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]],      dtype=np.intaa)
           _lowercase       :							List[Any]      		=		shift_tokens_right(lowerCamelCase,      1,      2)
           _lowercase       :							List[Any]      		=		np.equal(lowerCamelCase,      1).astype(np.floataa).sum()
           _lowercase       :							List[Any]      		=		np.equal(lowerCamelCase,      1).astype(np.floataa).sum()
           self.assertEqual(shifted.shape,      input_ids.shape)
           self.assertEqual(lowerCamelCase,      n_pad_before - 1)
           self.assertTrue(np.equal(shifted[:, 0],      2).all())






@require_flax
class _lowerCamelCase(						_a, unittest.TestCase, _a   ):
       lowercase_ :	int     =			True
       lowercase_ :	Optional[Any]     =			(
           (
               FlaxBlenderbotModel,
               FlaxBlenderbotForConditionalGeneration,
           )
           if is_flax_available()
           else ()
       )
       lowercase_ :	List[Any]     =			(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							int      		=		FlaxBlenderbotModelTester(self)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Tuple      		=		self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
               self.model_tester.check_use_cache_forward(lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Dict      		=		self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
               self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Tuple      		=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__):
                   _lowercase       :							Any      		=		self._prepare_for_class(lowerCamelCase,      lowerCamelCase)
                   _lowercase       :							str      		=		model_class(lowerCamelCase)

                   @jax.jit
                   def encode_jitted(lowerCamelCase,      lowerCamelCase=None,      **lowerCamelCase):
                       return model.encode(input_ids=lowerCamelCase,      attention_mask=lowerCamelCase)

                   with self.subTest('JIT Enabled'):
                       _lowercase       :							Any      		=		encode_jitted(**lowerCamelCase).to_tuple()

                   with self.subTest('JIT Disabled'):
                       with jax.disable_jit():
                           _lowercase       :							List[Any]      		=		encode_jitted(**lowerCamelCase).to_tuple()

                   self.assertEqual(len(lowerCamelCase),      len(lowerCamelCase))
                   for jitted_output, output in zip(lowerCamelCase,      lowerCamelCase):
                       self.assertEqual(jitted_output.shape,      output.shape)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Union[str, Any]      		=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__):
                   _lowercase       :							List[str]      		=		model_class(lowerCamelCase)
                   _lowercase       :							Optional[Any]      		=		model.encode(inputs_dict['input_ids'],      inputs_dict['attention_mask'])

                   _lowercase       :							str      		=		{
                       'decoder_input_ids': inputs_dict['decoder_input_ids'],
                       'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
                       'encoder_outputs': encoder_outputs,
                   }

                   @jax.jit
                   def decode_jitted(lowerCamelCase,      lowerCamelCase,      lowerCamelCase):
                       return model.decode(
                           decoder_input_ids=lowerCamelCase,      decoder_attention_mask=lowerCamelCase,      encoder_outputs=lowerCamelCase,      )

                   with self.subTest('JIT Enabled'):
                       _lowercase       :							List[Any]      		=		decode_jitted(**lowerCamelCase).to_tuple()

                   with self.subTest('JIT Disabled'):
                       with jax.disable_jit():
                           _lowercase       :							List[Any]      		=		decode_jitted(**lowerCamelCase).to_tuple()

                   self.assertEqual(len(lowerCamelCase),      len(lowerCamelCase))
                   for jitted_output, output in zip(lowerCamelCase,      lowerCamelCase):
                       self.assertEqual(jitted_output.shape,      output.shape)
       @slow
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           for model_class_name in self.all_model_classes:
               _lowercase       :							Optional[int]      		=		model_class_name.from_pretrained('facebook/blenderbot-400M-distill')
               # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
               _lowercase       :							Optional[Any]      		=		np.ones((1, 1)) * model.config.eos_token_id
               _lowercase       :							int      		=		model(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)




       @unittest.skipUnless(jax_device != 'cpu',      '3B test too slow on CPU.')
       @slow
       def 			UpperCamelCase     (					self)       ->   List[Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		{'num_beams': 1, 'early_stopping': True, 'min_length': 15, 'max_length': 25}
           _lowercase       :							List[Any]      		=		{'skip_special_tokens': True, 'clean_up_tokenization_spaces': True}

           _lowercase       :							List[str]      		=		FlaxBlenderbotForConditionalGeneration.from_pretrained('facebook/blenderbot-3B',      from_pt=lowerCamelCase)
           _lowercase       :							Dict      		=		BlenderbotTokenizer.from_pretrained('facebook/blenderbot-3B')

           _lowercase       :							Any      		=		['Sam']
           _lowercase       :							List[str]      		=		tokenizer(lowerCamelCase,      return_tensors='jax')

           _lowercase       :							Any      		=		model.generate(**lowerCamelCase,      **lowerCamelCase)
           _lowercase       :							str      		=		'Sam is a great name. It means "sun" in Gaelic.'

           _lowercase       :							List[str]      		=		tokenizer.batch_decode(lowerCamelCase,      **lowerCamelCase)
           assert generated_txt[0].strip() == tgt_text






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 _lowerCamelCase(						_a   ):
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Tuple      		=		self.config_class(**self.inputs_dict)
           self.parent.assertTrue(hasattr(lowerCamelCase,      'width_multiplier'))






class _lowerCamelCase:
       def __init__(					self,      lowerCamelCase,      lowerCamelCase=13,      lowerCamelCase=64,      lowerCamelCase=2,      lowerCamelCase=3,      lowerCamelCase="swish",      lowerCamelCase=3,      lowerCamelCase=32,      lowerCamelCase=0.1,      lowerCamelCase=0.0_2,      lowerCamelCase=True,      lowerCamelCase=True,      lowerCamelCase=10,      lowerCamelCase=None,      lowerCamelCase=0.2_5,      lowerCamelCase=0.0,      lowerCamelCase=0.0,      )       ->   Any:




           """simple docstring"""

           _lowercase       :							Any      		=		parent
           _lowercase       :							Optional[int]      		=		batch_size
           _lowercase       :							Dict      		=		image_size
           _lowercase       :							str      		=		patch_size
           _lowercase       :							Optional[int]      		=		num_channels
           _lowercase       :							Optional[Any]      		=		make_divisible(5_12 * width_multiplier,      divisor=8)
           _lowercase       :							str      		=		hidden_act
           _lowercase       :							Dict      		=		conv_kernel_size
           _lowercase       :							int      		=		output_stride
           _lowercase       :							Optional[Any]      		=		classifier_dropout_prob
           _lowercase       :							Tuple      		=		use_labels
           _lowercase       :							int      		=		is_training
           _lowercase       :							Optional[Any]      		=		num_labels
           _lowercase       :							Dict      		=		initializer_range
           _lowercase       :							List[str]      		=		scope
           _lowercase       :							Tuple      		=		width_multiplier
           _lowercase       :							List[str]      		=		ffn_dropout
           _lowercase       :							Dict      		=		attn_dropout
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Dict      		=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

           _lowercase       :							Dict      		=		None
           _lowercase       :							Optional[int]      		=		None
           if self.use_labels:
               _lowercase       :							Optional[Any]      		=		ids_tensor([self.batch_size],      self.num_labels)
               _lowercase       :							str      		=		ids_tensor([self.batch_size, self.image_size, self.image_size],      self.num_labels)

           _lowercase       :							Union[str, Any]      		=		self.get_config()

           return config, pixel_values, labels, pixel_labels
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """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,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Any:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		MobileViTVaModel(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							Optional[int]      		=		model(lowerCamelCase)
           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,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Optional[int]:




           """simple docstring"""

           _lowercase       :							int      		=		self.num_labels
           _lowercase       :							Optional[int]      		=		MobileViTVaForImageClassification(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							Optional[Any]      		=		model(lowerCamelCase,      labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_labels))
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   int:




           """simple docstring"""

           _lowercase       :							Any      		=		self.num_labels
           _lowercase       :							Union[str, Any]      		=		MobileViTVaForSemanticSegmentation(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							Optional[int]      		=		model(lowerCamelCase)
           self.parent.assertEqual(
               result.logits.shape,      (
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               ),      )
           _lowercase       :							List[Any]      		=		model(lowerCamelCase,      labels=lowerCamelCase)
           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"""

           _lowercase       :							str      		=		self.prepare_config_and_inputs()
           _lowercase		,				_lowercase		,				_lowercase		,				_lowercase       :							int      		=		config_and_inputs
           _lowercase       :							List[str]      		=		{'pixel_values': pixel_values}
           return config, inputs_dict






@require_torch
class _lowerCamelCase(						_a, _a, unittest.TestCase   ):
       lowercase_ :	List[Any]     =			(
           (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
           if is_torch_available()
           else ()
       )

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

       lowercase_ :	List[Any]     =			False
       lowercase_ :	Optional[int]     =			False
       lowercase_ :	List[Any]     =			False
       lowercase_ :	Tuple     =			False
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		MobileViTVaModelTester(self)
           _lowercase       :							Tuple      		=		MobileViTVaConfigTester(self,      config_class=lowerCamelCase,      has_text_modality=lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           self.config_tester.run_common_tests()
       @unittest.skip(reason='MobileViTV2 does not use inputs_embeds')
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

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




           """simple docstring"""

           pass
       @unittest.skip(reason='MobileViTV2 does not output attentions')
       def 			UpperCamelCase     (					self)       ->   List[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)       ->   int:




           """simple docstring"""

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




           """simple docstring"""

           pass
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase		,				_lowercase       :							Dict      		=		self.model_tester.prepare_config_and_inputs_for_common()

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

               _lowercase       :							Union[str, Any]      		=		['pixel_values']
               self.assertListEqual(arg_names[:1],      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           def check_hidden_states_output(lowerCamelCase,      lowerCamelCase,      lowerCamelCase):
               _lowercase       :							Optional[Any]      		=		model_class(lowerCamelCase)
               model.to(lowerCamelCase)
               model.eval()

               with torch.no_grad():
                   _lowercase       :							Optional[int]      		=		model(**self._prepare_for_class(lowerCamelCase,      lowerCamelCase))

               _lowercase       :							List[Any]      		=		outputs.hidden_states

               _lowercase       :							Tuple      		=		5
               self.assertEqual(len(lowerCamelCase),      lowerCamelCase)

               # MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
               # with the width and height being successively divided by 2.
               _lowercase       :							Optional[int]      		=		2
               for i in range(len(lowerCamelCase)):
                   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       :							Dict      		=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               _lowercase       :							Tuple      		=		True
               check_hidden_states_output(lowerCamelCase,      lowerCamelCase,      lowerCamelCase)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               _lowercase       :							Optional[Any]      		=		True

               check_hidden_states_output(lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							int      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							str      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase)




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




           """simple docstring"""

           for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase       :							str      		=		MobileViTVaModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)



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






@require_torch
@require_vision
class _lowerCamelCase(						unittest.TestCase   ):
       @cached_property
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """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"""

           _lowercase       :							List[str]      		=		MobileViTVaForImageClassification.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256').to(
               lowerCamelCase)

           _lowercase       :							Dict      		=		self.default_image_processor
           _lowercase       :							Union[str, Any]      		=		prepare_img()
           _lowercase       :							Dict      		=		image_processor(images=lowerCamelCase,      return_tensors='pt').to(lowerCamelCase)

           # forward pass
           with torch.no_grad():
               _lowercase       :							Tuple      		=		model(**lowerCamelCase)

           # verify the logits
           _lowercase       :							Optional[int]      		=		torch.Size((1, 10_00))
           self.assertEqual(outputs.logits.shape,      lowerCamelCase)

           _lowercase       :							Union[str, Any]      		=		torch.tensor([-1.63_36E00, -7.32_04E-02, -5.18_83E-01]).to(lowerCamelCase)

           self.assertTrue(torch.allclose(outputs.logits[0, :3],      lowerCamelCase,      atol=1E-4))
       @slow
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3')
           _lowercase       :							Optional[int]      		=		model.to(lowerCamelCase)

           _lowercase       :							Optional[int]      		=		MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3')

           _lowercase       :							Union[str, Any]      		=		prepare_img()
           _lowercase       :							Tuple      		=		image_processor(images=lowerCamelCase,      return_tensors='pt').to(lowerCamelCase)

           # forward pass
           with torch.no_grad():
               _lowercase       :							List[Any]      		=		model(**lowerCamelCase)
           _lowercase       :							str      		=		outputs.logits

           # verify the logits
           _lowercase       :							Tuple      		=		torch.Size((1, 21, 32, 32))
           self.assertEqual(logits.shape,      lowerCamelCase)

           _lowercase       :							Union[str, Any]      		=		torch.tensor(
               [
                   [[7.0_8_6_3, 7.1_5_2_5, 6.8_2_0_1], [6.6_9_3_1, 6.8_7_7_0, 6.8_9_3_3], [6.2_9_7_8, 7.0_3_6_6, 6.9_6_3_6]],
                   [[-3.7_1_3_4, -3.6_7_1_2, -3.6_6_7_5], [-3.5_8_2_5, -3.3_5_4_9, -3.4_7_7_7], [-3.3_4_3_5, -3.3_9_7_9, -3.2_8_5_7]],
                   [[-2.9_3_2_9, -2.8_0_0_3, -2.7_3_6_9], [-3.0_5_6_4, -2.4_7_8_0, -2.0_2_0_7], [-2.6_8_8_9, -1.9_2_9_8, -1.7_6_4_0]],
               ],      device=lowerCamelCase,      )

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




       @slow
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							List[str]      		=		MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3')
           _lowercase       :							Tuple      		=		model.to(lowerCamelCase)

           _lowercase       :							str      		=		MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3')

           _lowercase       :							int      		=		prepare_img()
           _lowercase       :							Dict      		=		image_processor(images=lowerCamelCase,      return_tensors='pt').to(lowerCamelCase)

           # forward pass
           with torch.no_grad():
               _lowercase       :							Union[str, Any]      		=		model(**lowerCamelCase)

           _lowercase       :							Any      		=		outputs.logits.detach().cpu()

           _lowercase       :							Optional[int]      		=		image_processor.post_process_semantic_segmentation(outputs=lowerCamelCase,      target_sizes=[(50, 60)])
           _lowercase       :							Any      		=		torch.Size((50, 60))
           self.assertEqual(segmentation[0].shape,      lowerCamelCase)

           _lowercase       :							Optional[Any]      		=		image_processor.post_process_semantic_segmentation(outputs=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		torch.Size((32, 32))
           self.assertEqual(segmentation[0].shape,      lowerCamelCase)


