





import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


_UpperCAmelCase										=					logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

_UpperCAmelCase										=					{
    'base': AutoModel,
    'sequence-classification': AutoModelForSequenceClassification,
    'question-answering': AutoModelForQuestionAnswering,
    'pretraining': AutoModelForPreTraining,
    'token-classification': AutoModelForTokenClassification,
    'language-modeling': AutoModelWithLMHead,
    'summarization': AutoModelForSeqaSeqLM,
    'translation': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
_UpperCAmelCase										=					{
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
_UpperCAmelCase										=					sorted(arg_to_scheduler.keys())
_UpperCAmelCase										=					'{' + ', '.join(arg_to_scheduler_choices) + '}'



class 				_UpperCamelCase       (						pl.LightningModule  ):

       def __init__(		self:			Dict       , _SCREAMING_SNAKE_CASE:			argparse.Namespace       , _SCREAMING_SNAKE_CASE:			Optional[int]=None       , _SCREAMING_SNAKE_CASE:			Dict="base"       , _SCREAMING_SNAKE_CASE:			Optional[int]=None       , _SCREAMING_SNAKE_CASE:			Any=None       , _SCREAMING_SNAKE_CASE:			Dict=None       , **_SCREAMING_SNAKE_CASE:			Tuple       , )							->			List[str]:




              """simple docstring"""


              super().__init__()
              # TODO: move to self.save_hyperparameters()
              # self.save_hyperparameters()
              # can also expand arguments into trainer signature for easier reading

              self.save_hyperparameters(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	0
              UpperCamelCase_ 							=	Path(self.hparams.output_dir		)
              UpperCamelCase_ 							=	self.hparams.cache_dir if self.hparams.cache_dir else None
              if config is None:
                     UpperCamelCase_ 							=	AutoConfig.from_pretrained(
                         self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path       , **({"num_labels": num_labels} if num_labels is not None else {})       , cache_dir=_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE       , )
              else:
                     UpperCamelCase_ 							=	config

              UpperCamelCase_ 							=	("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
              for p in extra_model_params:
                     if getattr(self.hparams       , _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE		):
                            assert hasattr(self.config       , _SCREAMING_SNAKE_CASE		), f'''model config doesn\'t have a `{p}` attribute'''
                            setattr(self.config       , _SCREAMING_SNAKE_CASE       , getattr(self.hparams       , _SCREAMING_SNAKE_CASE		)		)

              if tokenizer is None:
                     UpperCamelCase_ 							=	AutoTokenizer.from_pretrained(
                         self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path       , cache_dir=_SCREAMING_SNAKE_CASE       , )
              else:
                     UpperCamelCase_ 							=	tokenizer
              UpperCamelCase_ 							=	MODEL_MODES[mode]
              if model is None:
                     UpperCamelCase_ 							=	self.model_type.from_pretrained(
                         self.hparams.model_name_or_path       , from_tf=bool(".ckpt" in self.hparams.model_name_or_path		)       , config=self.config       , cache_dir=_SCREAMING_SNAKE_CASE       , )
              else:
                     UpperCamelCase_ 							=	model

       def lowercase			(		self:			Any       , *_SCREAMING_SNAKE_CASE:			Optional[Any]       , **_SCREAMING_SNAKE_CASE:			Union[str, Any]		)							->			Tuple:




              """simple docstring"""


              UpperCamelCase_ 							=	self.model_type.from_pretrained(*_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			Dict		)							->			Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	arg_to_scheduler[self.hparams.lr_scheduler]
              UpperCamelCase_ 							=	get_schedule_func(
                  self.opt       , num_warmup_steps=self.hparams.warmup_steps       , num_training_steps=self.total_steps()		)
              UpperCamelCase_ 							=	{"scheduler": scheduler, "interval": "step", "frequency": 1}
              return scheduler

       def lowercase			(		self:			List[str]		)							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	self.model
              UpperCamelCase_ 							=	["bias", "LayerNorm.weight"]
              UpperCamelCase_ 							=	[
                  {
                      "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay		)],
                      "weight_decay": self.hparams.weight_decay,
                  },
                  {
                      "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay		)],
                      "weight_decay": 0.0,
                  },
              ]
              if self.hparams.adafactor:
                     UpperCamelCase_ 							=	Adafactor(
                         _SCREAMING_SNAKE_CASE       , lr=self.hparams.learning_rate       , scale_parameter=_SCREAMING_SNAKE_CASE       , relative_step=_SCREAMING_SNAKE_CASE		)

              else:
                     UpperCamelCase_ 							=	AdamW(
                         _SCREAMING_SNAKE_CASE       , lr=self.hparams.learning_rate       , eps=self.hparams.adam_epsilon		)
              UpperCamelCase_ 							=	optimizer

              UpperCamelCase_ 							=	self.get_lr_scheduler()

              return [optimizer], [scheduler]

       def lowercase			(		self:			Dict       , _SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Optional[int]		)							->			Union[str, Any]:




              """simple docstring"""


              return self.validation_step(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Optional[int]		)							->			List[Any]:




              """simple docstring"""


              return self.validation_end(_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			List[Any]		)							->			int:




              """simple docstring"""


              UpperCamelCase_ 							=	max(1       , self.hparams.gpus		)  # TODO: consider num_tpu_cores
              UpperCamelCase_ 							=	self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
              return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs

       def lowercase			(		self:			Tuple       , _SCREAMING_SNAKE_CASE:			Union[str, Any]		)							->			List[str]:




              """simple docstring"""


              if mode == "test":
                     UpperCamelCase_ 							=	len(self.test_dataloader().dataset		)
              else:
                     UpperCamelCase_ 							=	self.get_dataloader("train"       , self.hparams.train_batch_size       , shuffle=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_ 							=	len(self.train_dataloader().dataset		)

       def lowercase			(		self:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			str       , _SCREAMING_SNAKE_CASE:			int       , _SCREAMING_SNAKE_CASE:			bool = False		)							->			int:




              """simple docstring"""


              raise NotImplementedError("You must implement this for your task"		)

       def lowercase			(		self:			Union[str, Any]		)							->			Dict:




              """simple docstring"""


              return self.train_loader

       def lowercase			(		self:			Tuple		)							->			Optional[Any]:




              """simple docstring"""


              return self.get_dataloader("dev"       , self.hparams.eval_batch_size       , shuffle=_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			Dict		)							->			List[Any]:




              """simple docstring"""


              return self.get_dataloader("test"       , self.hparams.eval_batch_size       , shuffle=_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			Optional[int]       , _SCREAMING_SNAKE_CASE:			List[Any]		)							->			Dict:




              """simple docstring"""


              return os.path.join(
                  self.hparams.data_dir       , "cached_{}_{}_{}".format(
                      _SCREAMING_SNAKE_CASE       , list(filter(_SCREAMING_SNAKE_CASE       , self.hparams.model_name_or_path.split("/"		)		)		).pop()       , str(self.hparams.max_seq_length		)       , )       , )

       @pl.utilities.rank_zero_only
       def lowercase			(		self:			str       , _SCREAMING_SNAKE_CASE:			Dict[str, Any]		)							->			None:




              """simple docstring"""


              UpperCamelCase_ 							=	self.output_dir.joinpath("best_tfmr"		)
              UpperCamelCase_ 							=	self.step_count
              self.model.save_pretrained(_SCREAMING_SNAKE_CASE		)
              self.tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE		)



       @staticmethod
       def lowercase			(		_SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			List[Any]		)							->			Tuple:




              """simple docstring"""


              parser.add_argument(
                  "--model_name_or_path"       , default=_SCREAMING_SNAKE_CASE       , type=_SCREAMING_SNAKE_CASE       , required=_SCREAMING_SNAKE_CASE       , help="Path to pretrained model or model identifier from huggingface.co/models"       , )
              parser.add_argument(
                  "--config_name"       , default=""       , type=_SCREAMING_SNAKE_CASE       , help="Pretrained config name or path if not the same as model_name"		)
              parser.add_argument(
                  "--tokenizer_name"       , default=_SCREAMING_SNAKE_CASE       , type=_SCREAMING_SNAKE_CASE       , help="Pretrained tokenizer name or path if not the same as model_name"       , )
              parser.add_argument(
                  "--cache_dir"       , default=""       , type=_SCREAMING_SNAKE_CASE       , help="Where do you want to store the pre-trained models downloaded from huggingface.co"       , )
              parser.add_argument(
                  "--encoder_layerdrop"       , type=_SCREAMING_SNAKE_CASE       , help="Encoder layer dropout probability (Optional). Goes into model.config"       , )
              parser.add_argument(
                  "--decoder_layerdrop"       , type=_SCREAMING_SNAKE_CASE       , help="Decoder layer dropout probability (Optional). Goes into model.config"       , )
              parser.add_argument(
                  "--dropout"       , type=_SCREAMING_SNAKE_CASE       , help="Dropout probability (Optional). Goes into model.config"       , )
              parser.add_argument(
                  "--attention_dropout"       , type=_SCREAMING_SNAKE_CASE       , help="Attention dropout probability (Optional). Goes into model.config"       , )
              parser.add_argument("--learning_rate"       , default=5e-5       , type=_SCREAMING_SNAKE_CASE       , help="The initial learning rate for Adam."		)
              parser.add_argument(
                  "--lr_scheduler"       , default="linear"       , choices=_SCREAMING_SNAKE_CASE       , metavar=_SCREAMING_SNAKE_CASE       , type=_SCREAMING_SNAKE_CASE       , help="Learning rate scheduler"       , )
              parser.add_argument("--weight_decay"       , default=0.0       , type=_SCREAMING_SNAKE_CASE       , help="Weight decay if we apply some."		)
              parser.add_argument("--adam_epsilon"       , default=1e-8       , type=_SCREAMING_SNAKE_CASE       , help="Epsilon for Adam optimizer."		)
              parser.add_argument("--warmup_steps"       , default=0       , type=_SCREAMING_SNAKE_CASE       , help="Linear warmup over warmup_steps."		)
              parser.add_argument("--num_workers"       , default=4       , type=_SCREAMING_SNAKE_CASE       , help="kwarg passed to DataLoader"		)
              parser.add_argument("--num_train_epochs"       , dest="max_epochs"       , default=3       , type=_SCREAMING_SNAKE_CASE		)
              parser.add_argument("--train_batch_size"       , default=32       , type=_SCREAMING_SNAKE_CASE		)
              parser.add_argument("--eval_batch_size"       , default=32       , type=_SCREAMING_SNAKE_CASE		)
              parser.add_argument("--adafactor"       , action="store_true"		)



class 				_UpperCamelCase       (						pl.Callback  ):

       def lowercase			(		self:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			str       , _SCREAMING_SNAKE_CASE:			int		)							->			List[str]:




              """simple docstring"""


              UpperCamelCase_ 							=	trainer.lr_schedulers[0]["scheduler"]
              UpperCamelCase_ 							=	{f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()		)}
              pl_module.logger.log_metrics(_SCREAMING_SNAKE_CASE		)

       def lowercase			(		self:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			pl.Trainer       , _SCREAMING_SNAKE_CASE:			pl.LightningModule		)							->			Tuple:




              """simple docstring"""


              rank_zero_info("***** Validation results *****"		)
              UpperCamelCase_ 							=	trainer.callback_metrics
              # Log results
              for key in sorted(_SCREAMING_SNAKE_CASE		):
                     if key not in ["log", "progress_bar"]:
                            rank_zero_info("{} = {}\n".format(_SCREAMING_SNAKE_CASE       , str(metrics[key]		)		)		)



       def lowercase			(		self:			Tuple       , _SCREAMING_SNAKE_CASE:			pl.Trainer       , _SCREAMING_SNAKE_CASE:			pl.LightningModule		)							->			int:




              """simple docstring"""


              rank_zero_info("***** Test results *****"		)
              UpperCamelCase_ 							=	trainer.callback_metrics
              # Log and save results to file
              UpperCamelCase_ 							=	os.path.join(pl_module.hparams.output_dir       , "test_results.txt"		)
              with open(_SCREAMING_SNAKE_CASE       , "w"		) as writer:
                     for key in sorted(_SCREAMING_SNAKE_CASE		):
                            if key not in ["log", "progress_bar"]:
                                   rank_zero_info("{} = {}\n".format(_SCREAMING_SNAKE_CASE       , str(metrics[key]		)		)		)
                                   writer.write("{} = {}\n".format(_SCREAMING_SNAKE_CASE       , str(metrics[key]		)		)		)




def 			lowerCAmelCase_			(				UpperCamelCase_       ,		UpperCamelCase_	)  ->		None:
       #  To allow all pl args uncomment the following line
       #  parser = pl.Trainer.add_argparse_args(parser)
       parser.add_argument(
           "--output_dir"       ,		default=UpperCamelCase_       ,		type=UpperCamelCase_       ,		required=UpperCamelCase_       ,		help="The output directory where the model predictions and checkpoints will be written."       ,		)
       parser.add_argument(
           "--fp16"       ,		action="store_true"       ,		help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"       ,		)

       parser.add_argument(
           "--fp16_opt_level"       ,		type=UpperCamelCase_       ,		default="O2"       ,		help=(
               "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
               "See details at https://nvidia.github.io/apex/amp.html"
           )       ,		)
       parser.add_argument("--n_tpu_cores"       ,		dest="tpu_cores"       ,		type=UpperCamelCase_	)
       parser.add_argument("--max_grad_norm"       ,		dest="gradient_clip_val"       ,		default=1.0       ,		type=UpperCamelCase_       ,		help="Max gradient norm"	)
       parser.add_argument("--do_train"       ,		action="store_true"       ,		help="Whether to run training."	)
       parser.add_argument("--do_predict"       ,		action="store_true"       ,		help="Whether to run predictions on the test set."	)
       parser.add_argument(
           "--gradient_accumulation_steps"       ,		dest="accumulate_grad_batches"       ,		type=UpperCamelCase_       ,		default=1       ,		help="Number of updates steps to accumulate before performing a backward/update pass."       ,		)
       parser.add_argument("--seed"       ,		type=UpperCamelCase_       ,		default=42       ,		help="random seed for initialization"	)
       parser.add_argument(
           "--data_dir"       ,		default=UpperCamelCase_       ,		type=UpperCamelCase_       ,		required=UpperCamelCase_       ,		help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."       ,		)







def 			lowerCAmelCase_			(				UpperCamelCase_       ,		UpperCamelCase_       ,		UpperCamelCase_=None       ,		UpperCamelCase_=True       ,		UpperCamelCase_=[]       ,		UpperCamelCase_=None       ,		UpperCamelCase_=None       ,		**UpperCamelCase_       ,		)  ->		Tuple:
       pl.seed_everything(args.seed	)

       # init model
       UpperCamelCase_ 							=	Path(model.hparams.output_dir	)
       odir.mkdir(exist_ok=UpperCamelCase_	)

       # add custom checkpoints
       if checkpoint_callback is None:
              UpperCamelCase_ 							=	pl.callbacks.ModelCheckpoint(
                  filepath=args.output_dir       ,		prefix="checkpoint"       ,		monitor="val_loss"       ,		mode="min"       ,		save_top_k=1	)
       if early_stopping_callback:
              extra_callbacks.append(UpperCamelCase_	)
       if logging_callback is None:
              UpperCamelCase_ 							=	LoggingCallback()

       UpperCamelCase_ 							=	{}

       # TODO: remove with PyTorch 1.6 since pl uses native amp
       if args.fpaa:
              UpperCamelCase_ 							=	16
              UpperCamelCase_ 							=	args.fpaa_opt_level

       if args.gpus > 1:
              UpperCamelCase_ 							=	"ddp"

       UpperCamelCase_ 							=	args.accumulate_grad_batches
       UpperCamelCase_ 							=	extra_train_kwargs.get("accelerator"       ,		UpperCamelCase_	)
       UpperCamelCase_ 							=	extra_train_kwargs.get("profiler"       ,		UpperCamelCase_	)

       UpperCamelCase_ 							=	pl.Trainer.from_argparse_args(
           UpperCamelCase_       ,		weights_summary=UpperCamelCase_       ,		callbacks=[logging_callback] + extra_callbacks       ,		logger=UpperCamelCase_       ,		checkpoint_callback=UpperCamelCase_       ,		**UpperCamelCase_       ,		)

       if args.do_train:
              trainer.fit(UpperCamelCase_	)

       return trainer





"""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 			SCREAMING_SNAKE_CASE_  (   _lowercase):




							'''simple docstring'''





							def        UpperCAmelCase (						self)   ->	Dict:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      self.config_class(**self.inputs_dict)
								self.parent.assertTrue(hasattr(lowerCamelCase__	,					"hidden_sizes"))
								self.parent.assertTrue(hasattr(lowerCamelCase__	,					"neck_hidden_sizes"))
								self.parent.assertTrue(hasattr(lowerCamelCase__	,					"num_attention_heads"))


class 			SCREAMING_SNAKE_CASE_  :




							'''simple docstring'''





							def __init__(						self	,					lowerCamelCase__	,					lowerCamelCase__=13	,					lowerCamelCase__=32	,					lowerCamelCase__=2	,					lowerCamelCase__=3	,					lowerCamelCase__=640	,					lowerCamelCase__=4	,					lowerCamelCase__="silu"	,					lowerCamelCase__=3	,					lowerCamelCase__=32	,					lowerCamelCase__=0.1	,					lowerCamelCase__=0.1	,					lowerCamelCase__=0.1	,					lowerCamelCase__=0.02	,					lowerCamelCase__=True	,					lowerCamelCase__=True	,					lowerCamelCase__=10	,					lowerCamelCase__=None	,					)   ->	Dict:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      parent
								snake_case__  :						List[Any]		=      batch_size
								snake_case__  :						Any		=      image_size
								snake_case__  :						Tuple		=      patch_size
								snake_case__  :						int		=      num_channels
								snake_case__  :						Optional[Any]		=      last_hidden_size
								snake_case__  :						Union[str, Any]		=      num_attention_heads
								snake_case__  :						Tuple		=      hidden_act
								snake_case__  :						Optional[int]		=      conv_kernel_size
								snake_case__  :						int		=      output_stride
								snake_case__  :						Optional[Any]		=      hidden_dropout_prob
								snake_case__  :						str		=      attention_probs_dropout_prob
								snake_case__  :						Any		=      classifier_dropout_prob
								snake_case__  :						Tuple		=      use_labels
								snake_case__  :						Optional[Any]		=      is_training
								snake_case__  :						Dict		=      num_labels
								snake_case__  :						Dict		=      initializer_range
								snake_case__  :						List[str]		=      scope




							def        UpperCAmelCase (						self)   ->	Optional[Any]:







								'''simple docstring'''




								snake_case__  :						List[Any]		=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

								snake_case__  :						Dict		=      self.get_config()

								return config, pixel_values, labels, pixel_labels




							def        UpperCAmelCase (						self)   ->	str:







								'''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	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)   ->	Optional[int]:







								'''simple docstring'''




								snake_case__  :						Tuple		=      TFMobileViTModel(config=lowerCamelCase__)
								snake_case__  :						int		=      model(lowerCamelCase__	,					training=lowerCamelCase__)
								snake_case__  :						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        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)   ->	Optional[int]:







								'''simple docstring'''




								snake_case__  :						Any		=      self.num_labels
								snake_case__  :						Dict		=      TFMobileViTForImageClassification(lowerCamelCase__)
								snake_case__  :						List[Any]		=      model(lowerCamelCase__	,					labels=lowerCamelCase__	,					training=lowerCamelCase__)
								self.parent.assertEqual(result.logits.shape	,					(self.batch_size, self.num_labels))




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)   ->	str:







								'''simple docstring'''




								snake_case__  :						Optional[Any]		=      self.num_labels
								snake_case__  :						Any		=      TFMobileViTForSemanticSegmentation(lowerCamelCase__)
								snake_case__  :						List[Any]		=      self.image_size // self.output_stride

								snake_case__  :						Tuple		=      model(lowerCamelCase__	,					training=lowerCamelCase__)
								self.parent.assertEqual(
								    result.logits.shape	,					(self.batch_size, self.num_labels, expected_height, expected_width))

								snake_case__  :						str		=      model(lowerCamelCase__	,					labels=lowerCamelCase__	,					training=lowerCamelCase__)
								self.parent.assertEqual(
								    result.logits.shape	,					(self.batch_size, self.num_labels, expected_height, expected_width))







							def        UpperCAmelCase (						self)   ->	Tuple:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      self.prepare_config_and_inputs()
								snake_case__,					snake_case__,					snake_case__,					snake_case__  :						List[Any]		=      config_and_inputs
								snake_case__  :						Tuple		=      {"pixel_values": pixel_values}
								return config, inputs_dict


@require_tf
class 			SCREAMING_SNAKE_CASE_  (   _lowercase	,      _lowercase	,      unittest.TestCase):




							'''simple docstring'''





							__magic_name__					:     Optional[int]				  =  (
							    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
							    if is_tf_available()
							    else ()
							)
							__magic_name__					:     str				  =  (
							    {'''feature-extraction''': TFMobileViTModel, '''image-classification''': TFMobileViTForImageClassification}
							    if is_tf_available()
							    else {}
							)

							__magic_name__					:     List[str]				  =  False
							__magic_name__					:     int				  =  False
							__magic_name__					:     Any				  =  False
							__magic_name__					:     Optional[int]				  =  False
							__magic_name__					:     Union[str, Any]				  =  False




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







								'''simple docstring'''




								snake_case__  :						Any		=      TFMobileViTModelTester(self)
								snake_case__  :						Dict		=      TFMobileViTConfigTester(self	,					config_class=lowerCamelCase__	,					has_text_modality=lowerCamelCase__)




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







								'''simple docstring'''




								self.config_tester.run_common_tests()




							@unittest.skip(reason="MobileViT does not use inputs_embeds")
							def        UpperCAmelCase (						self)   ->	Optional[Any]:







								'''simple docstring'''




								pass




							@unittest.skip(reason="MobileViT does not support input and output embeddings")
							def        UpperCAmelCase (						self)   ->	List[Any]:







								'''simple docstring'''




								pass




							@unittest.skip(reason="MobileViT does not output attentions")
							def        UpperCAmelCase (						self)   ->	Dict:







								'''simple docstring'''




								pass




							def        UpperCAmelCase (						self)   ->	List[Any]:







								'''simple docstring'''




								snake_case__,					snake_case__  :						int		=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									snake_case__  :						str		=      model_class(lowerCamelCase__)
									snake_case__  :						Optional[int]		=      inspect.signature(model.call)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									snake_case__  :						List[Any]		=      [*signature.parameters.keys()]

									snake_case__  :						Dict		=      ["pixel_values"]
									self.assertListEqual(arg_names[:1]	,					lowerCamelCase__)




							def        UpperCAmelCase (						self)   ->	Optional[Any]:







								'''simple docstring'''




								snake_case__  :						Dict		=      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__):
									snake_case__  :						List[Any]		=      model_class(lowerCamelCase__)

									snake_case__  :						Optional[int]		=      model(**self._prepare_for_class(lowerCamelCase__	,					lowerCamelCase__))

									snake_case__  :						List[Any]		=      outputs.hidden_states

									snake_case__  :						Union[str, Any]		=      5
									self.assertEqual(len(lowerCamelCase__)	,					lowerCamelCase__)

									# 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__  :						str		=      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)

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

								for model_class in self.all_model_classes:
									snake_case__  :						List[Any]		=      True
									check_hidden_states_output(lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)

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

									check_hidden_states_output(lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__)




							def        UpperCAmelCase (						self)   ->	List[str]:







								'''simple docstring'''




								snake_case__  :						Dict		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowerCamelCase__)




							def        UpperCAmelCase (						self)   ->	str:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase__)




							@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        UpperCAmelCase (						self)   ->	Any:







								'''simple docstring'''




								super().test_dataset_conversion()




							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__=2E-1	,					lowerCamelCase__=2E-1)   ->	Optional[Any]:







								'''simple docstring'''




								self.assertTrue(np.allclose(lowerCamelCase__	,					lowerCamelCase__	,					atol=lowerCamelCase__	,					rtol=lowerCamelCase__))




							@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        UpperCAmelCase (						self)   ->	int:







								'''simple docstring'''




								snake_case__,					snake_case__  :						str		=      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":
										snake_case__  :						Dict		=      model_class(lowerCamelCase__)
										if getattr(lowerCamelCase__	,					"hf_compute_loss"	,					lowerCamelCase__):
											super().test_keras_fit()




							def        UpperCAmelCase (						self)   ->	str:







								'''simple docstring'''




								snake_case__,					snake_case__  :						List[str]		=      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":
										snake_case__  :						Any		=      5

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

										# Test that model correctly compute the loss with kwargs
										snake_case__  :						Dict		=      self._prepare_for_class(inputs_dict.copy()	,					lowerCamelCase__	,					return_labels=lowerCamelCase__)
										snake_case__  :						Dict		=      {"input_ids", "pixel_values", "input_features"}
										snake_case__  :						Any		=      possible_input_names.intersection(set(lowerCamelCase__)).pop()
										snake_case__  :						List[str]		=      prepared_for_class.pop(lowerCamelCase__)

										snake_case__  :						Optional[Any]		=      model(lowerCamelCase__	,					**lowerCamelCase__)[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
										snake_case__  :						Dict		=      self._prepare_for_class(inputs_dict.copy()	,					lowerCamelCase__	,					return_labels=lowerCamelCase__)
										snake_case__  :						Dict		=      {"input_ids", "pixel_values", "input_features"}
										snake_case__  :						List[Any]		=      possible_input_names.intersection(set(lowerCamelCase__)).pop()
										snake_case__  :						List[Any]		=      prepared_for_class.pop(lowerCamelCase__)
										if "labels" in prepared_for_class:
											snake_case__  :						int		=      prepared_for_class["labels"].numpy()
											if len(labels.shape) > 1 and labels.shape[1] != 1:
												# labels[0] = -100
												snake_case__  :						Union[str, Any]		=      tf.convert_to_tensor(lowerCamelCase__)
												snake_case__  :						Tuple		=      model(lowerCamelCase__	,					**lowerCamelCase__)[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
										snake_case__  :						int		=      self._prepare_for_class(inputs_dict.copy()	,					lowerCamelCase__	,					return_labels=lowerCamelCase__)
										snake_case__  :						Union[str, Any]		=      model(lowerCamelCase__)[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
										snake_case__  :						str		=      self._prepare_for_class(inputs_dict.copy()	,					lowerCamelCase__	,					return_labels=lowerCamelCase__)

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

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

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

										for index, value in sorted_tuple_index_mapping:
											snake_case__  :						Any		=      prepared_for_class[value]

										snake_case__  :						List[str]		=      tuple(lowerCamelCase__)

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

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







							@slow
							def        UpperCAmelCase (						self)   ->	str:







								'''simple docstring'''




								for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									snake_case__  :						Optional[Any]		=      TFMobileViTModel.from_pretrained(lowerCamelCase__)
									self.assertIsNotNone(lowerCamelCase__)



def 			A__		(					)							->	Optional[Any]:

	'''simple docstring'''


	snake_case__  :						Any		=      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
	return image



@require_tf
class 			SCREAMING_SNAKE_CASE_  (   unittest.TestCase):




							'''simple docstring'''





							@slow
							def        UpperCAmelCase (						self)   ->	str:







								'''simple docstring'''




								snake_case__  :						Union[str, Any]		=      TFMobileViTForImageClassification.from_pretrained("apple/mobilevit-xx-small")

								snake_case__  :						List[str]		=      MobileViTImageProcessor.from_pretrained("apple/mobilevit-xx-small")
								snake_case__  :						Dict		=      prepare_img()
								snake_case__  :						List[str]		=      image_processor(images=lowerCamelCase__	,					return_tensors="tf")

								# forward pass
								snake_case__  :						Union[str, Any]		=      model(**lowerCamelCase__	,					training=lowerCamelCase__)

								# verify the logits
								snake_case__  :						str		=      tf.TensorShape((1, 1_000))
								self.assertEqual(outputs.logits.shape	,					lowerCamelCase__)

								snake_case__  :						List[Any]		=      tf.constant([-1.93_64, -1.23_27, -0.46_53])

								tf.debugging.assert_near(outputs.logits[0, :3]	,					lowerCamelCase__	,					atol=1E-4	,					rtol=1E-04)







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







								'''simple docstring'''




								snake_case__  :						Dict		=      TFMobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small")

								snake_case__  :						Any		=      MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")

								snake_case__  :						List[str]		=      prepare_img()
								snake_case__  :						Dict		=      image_processor(images=lowerCamelCase__	,					return_tensors="tf")

								# forward pass
								snake_case__  :						str		=      model(inputs.pixel_values	,					training=lowerCamelCase__)
								snake_case__  :						Dict		=      outputs.logits

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

								snake_case__  :						List[Any]		=      tf.constant(
								    [
								        [[6.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
								        [[-10.68_69, -10.32_50, -10.34_71], [-10.42_28, -9.98_68, -9.71_32], [-11.04_05, -11.02_21, -10.73_18]],
								        [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
								    ])

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





