





import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
     from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import MobileViTImageProcessor



class 				_UpperCAmelCase    (    _UpperCamelCase    ):

  """simple docstring"""


  def 				lowercase		(				self :		Tuple     )							->						int:
     __lowerCAmelCase								=	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 				_UpperCAmelCase    :

  """simple docstring"""


  def __init__(				self :		Optional[Any]					,  lowerCAmelCase_ :		Union[str, Any]					,  lowerCAmelCase_ :		Any=1_3					,  lowerCAmelCase_ :		Dict=3_2					,  lowerCAmelCase_ :		Dict=2					,  lowerCAmelCase_ :		Optional[Any]=3					,  lowerCAmelCase_ :		List[str]=6_4_0					,  lowerCAmelCase_ :		Dict=4					,  lowerCAmelCase_ :		Union[str, Any]="silu"					,  lowerCAmelCase_ :		Any=3					,  lowerCAmelCase_ :		int=3_2					,  lowerCAmelCase_ :		Tuple=0.1					,  lowerCAmelCase_ :		int=0.1					,  lowerCAmelCase_ :		List[Any]=0.1					,  lowerCAmelCase_ :		List[Any]=0.02					,  lowerCAmelCase_ :		str=True					,  lowerCAmelCase_ :		List[str]=True					,  lowerCAmelCase_ :		Optional[Any]=1_0					,  lowerCAmelCase_ :		Any=None					,  )							->						Dict:
     __lowerCAmelCase								=	parent
     __lowerCAmelCase								=	batch_size
     __lowerCAmelCase								=	image_size
     __lowerCAmelCase								=	patch_size
     __lowerCAmelCase								=	num_channels
     __lowerCAmelCase								=	last_hidden_size
     __lowerCAmelCase								=	num_attention_heads
     __lowerCAmelCase								=	hidden_act
     __lowerCAmelCase								=	conv_kernel_size
     __lowerCAmelCase								=	output_stride
     __lowerCAmelCase								=	hidden_dropout_prob
     __lowerCAmelCase								=	attention_probs_dropout_prob
     __lowerCAmelCase								=	classifier_dropout_prob
     __lowerCAmelCase								=	use_labels
     __lowerCAmelCase								=	is_training
     __lowerCAmelCase								=	num_labels
     __lowerCAmelCase								=	initializer_range
     __lowerCAmelCase								=	scope






  def 				lowercase		(				self :		List[Any]     )							->						Any:
     __lowerCAmelCase								=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

     __lowerCAmelCase								=	self.get_config()

     return config, pixel_values, labels, pixel_labels






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






  def 				lowercase		(				self :		Optional[int]					,  lowerCAmelCase_ :		List[str]					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		Union[str, Any]     )							->						List[Any]:
     __lowerCAmelCase								=	self.num_labels
     __lowerCAmelCase								=	MobileViTForSemanticSegmentation(lowerCAmelCase_     )
     model.to(lowerCAmelCase_     )
     model.eval()
     __lowerCAmelCase								=	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,
         )					,  )
     __lowerCAmelCase								=	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 				lowercase		(				self :		int     )							->						Dict:
     __lowerCAmelCase								=	self.prepare_config_and_inputs()
     __lowerCAmelCase	,	__lowerCAmelCase	,	__lowerCAmelCase	,	__lowerCAmelCase								=	config_and_inputs
     __lowerCAmelCase								=	{'pixel_values': pixel_values}
     return config, inputs_dict





@require_torch
class 				_UpperCAmelCase    (    _UpperCamelCase		,    _UpperCamelCase		,    unittest.TestCase    ):

  """simple docstring"""


  a_        = (
      (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
      if is_torch_available()
      else ()
  )
  a_        = (
      {
          """feature-extraction""": MobileViTModel,
          """image-classification""": MobileViTForImageClassification,
          """image-segmentation""": MobileViTForSemanticSegmentation,
      }
      if is_torch_available()
      else {}
  )

  a_        = False
  a_        = False
  a_        = False
  a_        = False






  def 				lowercase		(				self :		List[str]     )							->						Union[str, Any]:
     __lowerCAmelCase								=	MobileViTModelTester(self     )
     __lowerCAmelCase								=	MobileViTConfigTester(self					,  config_class=lowerCAmelCase_					,  has_text_modality=lowerCAmelCase_     )






  def 				lowercase		(				self :		Optional[int]     )							->						int:
     self.config_tester.run_common_tests()






  @unittest.skip(reason='MobileViT does not use inputs_embeds'     )
  def 				lowercase		(				self :		List[str]     )							->						int:
     pass






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






  @unittest.skip(reason='MobileViT does not output attentions'     )
  def 				lowercase		(				self :		List[str]     )							->						Optional[Any]:
     pass






  def 				lowercase		(				self :		List[Any]     )							->						Dict:
     __lowerCAmelCase	,	__lowerCAmelCase								=	self.model_tester.prepare_config_and_inputs_for_common()

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

        __lowerCAmelCase								=	['pixel_values']
        self.assertListEqual(arg_names[:1]					,  lowerCAmelCase_     )






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






  def 				lowercase		(				self :		int     )							->						Optional[Any]:
     __lowerCAmelCase								=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowerCAmelCase_     )






  def 				lowercase		(				self :		Tuple     )							->						Union[str, Any]:
     def check_hidden_states_output(lowerCAmelCase_ :		Any					,  lowerCAmelCase_ :		List[str]					,  lowerCAmelCase_ :		Optional[int]     ):
        __lowerCAmelCase								=	model_class(lowerCAmelCase_     )
        model.to(lowerCAmelCase_     )
        model.eval()

        with torch.no_grad():
           __lowerCAmelCase								=	model(**self._prepare_for_class(lowerCAmelCase_					,  lowerCAmelCase_     )     )

        __lowerCAmelCase								=	outputs.hidden_states

        __lowerCAmelCase								=	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.
        __lowerCAmelCase								=	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     )

     __lowerCAmelCase	,	__lowerCAmelCase								=	self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        __lowerCAmelCase								=	True
        check_hidden_states_output(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )

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

        check_hidden_states_output(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )






  def 				lowercase		(				self :		int     )							->						Dict:
     __lowerCAmelCase								=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase_     )






  def 				lowercase		(				self :		List[str]     )							->						List[str]:
     __lowerCAmelCase								=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_semantic_segmentation(*lowerCAmelCase_     )






  @slow
  def 				lowercase		(				self :		str     )							->						Optional[int]:
     for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase								=	MobileViTModel.from_pretrained(lowerCAmelCase_     )
        self.assertIsNotNone(lowerCAmelCase_     )





def   a_    ( ):
   __lowerCAmelCase								=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
   return image




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

  """simple docstring"""


  @cached_property
  def 				lowercase		(				self :		Dict     )							->						Optional[Any]:
     return MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'     ) if is_vision_available() else None






  @slow
  def 				lowercase		(				self :		Optional[int]     )							->						List[Any]:
     __lowerCAmelCase								=	MobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'     ).to(lowerCAmelCase_     )

     __lowerCAmelCase								=	self.default_image_processor
     __lowerCAmelCase								=	prepare_img()
     __lowerCAmelCase								=	image_processor(images=lowerCAmelCase_					,  return_tensors='pt'     ).to(lowerCAmelCase_     )

     # forward pass
     with torch.no_grad():
        __lowerCAmelCase								=	model(**lowerCAmelCase_     )

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

     __lowerCAmelCase								=	torch.tensor([-1.93_64, -1.23_27, -0.46_53]     ).to(lowerCAmelCase_     )

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






  @slow
  def 				lowercase		(				self :		List[Any]     )							->						int:
     __lowerCAmelCase								=	MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )
     __lowerCAmelCase								=	model.to(lowerCAmelCase_     )

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

     __lowerCAmelCase								=	prepare_img()
     __lowerCAmelCase								=	image_processor(images=lowerCAmelCase_					,  return_tensors='pt'     ).to(lowerCAmelCase_     )

     # forward pass
     with torch.no_grad():
        __lowerCAmelCase								=	model(**lowerCAmelCase_     )
     __lowerCAmelCase								=	outputs.logits

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

     __lowerCAmelCase								=	torch.tensor(
         [
             [[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]],
         ]					,  device=lowerCAmelCase_					,  )

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






  @slow
  def 				lowercase		(				self :		Optional[Any]     )							->						Any:
     __lowerCAmelCase								=	MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'     )
     __lowerCAmelCase								=	model.to(lowerCAmelCase_     )

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

     __lowerCAmelCase								=	prepare_img()
     __lowerCAmelCase								=	image_processor(images=lowerCAmelCase_					,  return_tensors='pt'     ).to(lowerCAmelCase_     )

     # forward pass
     with torch.no_grad():
        __lowerCAmelCase								=	model(**lowerCAmelCase_     )

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

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

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








import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json

def   a_    ( lowerCAmelCase_      :			Optional[int]	):
   __lowerCAmelCase								=	filter(lambda lowerCAmelCase_	: p.requires_grad,	model.parameters()	)
   __lowerCAmelCase								=	sum([np.prod(p.size()	) for p in model_parameters]	)
   return params


_snake_case  :    Dict						    =  logging.getLogger(__name__)

def   a_    ( lowerCAmelCase_      :			Optional[int],	lowerCAmelCase_      :			Optional[int]	):
   if metric == "rouge2":
      __lowerCAmelCase								=	'{val_avg_rouge2:.4f}-{step_count}'
   elif metric == "bleu":
      __lowerCAmelCase								=	'{val_avg_bleu:.4f}-{step_count}'
   elif metric == "em":
      __lowerCAmelCase								=	'{val_avg_em:.4f}-{step_count}'
   else:
      raise NotImplementedError(
          F"""seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this"""
          ' function.'	)

   __lowerCAmelCase								=	ModelCheckpoint(
       dirpath=lowerCAmelCase_,	filename=lowerCAmelCase_,	monitor=F"""val_{metric}""",	mode='max',	save_top_k=3,	every_n_epochs=1,	)
   return checkpoint_callback

def   a_    ( lowerCAmelCase_      :			Union[str, Any],	lowerCAmelCase_      :			Any	):
   return EarlyStopping(
       monitor=F"""val_{metric}""",	mode='min' if 'loss' in metric else 'max',	patience=lowerCAmelCase_,	verbose=lowerCAmelCase_,	)




class 				_UpperCAmelCase    (    pl.Callback    ):

  """simple docstring"""


  def 				lowercase		(				self :		Tuple					,  lowerCAmelCase_ :		List[str]					,  lowerCAmelCase_ :		int     )							->						Any:
     __lowerCAmelCase								=	{f"""lr_group_{i}""": param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups     )}
     pl_module.logger.log_metrics(lowerCAmelCase_     )






  @rank_zero_only
  def 				lowercase		(				self :		Optional[int]					,  lowerCAmelCase_ :		pl.Trainer					,  lowerCAmelCase_ :		pl.LightningModule					,  lowerCAmelCase_ :		str					,  lowerCAmelCase_ :		List[Any]=True     )							->						None:
     logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""     )
     __lowerCAmelCase								=	trainer.callback_metrics
     trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}     )
     # Log results
     __lowerCAmelCase								=	Path(pl_module.hparams.output_dir     )
     if type_path == "test":
        __lowerCAmelCase								=	od / 'test_results.txt'
        __lowerCAmelCase								=	od / 'test_generations.txt'
     else:
        # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
        # If people want this it will be easy enough to add back.
        __lowerCAmelCase								=	od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
        __lowerCAmelCase								=	od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
        results_file.parent.mkdir(exist_ok=lowerCAmelCase_     )
        generations_file.parent.mkdir(exist_ok=lowerCAmelCase_     )
     with open(lowerCAmelCase_					,  'a+'     ) as writer:
        for key in sorted(lowerCAmelCase_     ):
           if key in ["log", "progress_bar", "preds"]:
              continue
           __lowerCAmelCase								=	metrics[key]
           if isinstance(lowerCAmelCase_					,  torch.Tensor     ):
              __lowerCAmelCase								=	val.item()
           __lowerCAmelCase								=	f"""{key}: {val:.6f}\n"""
           writer.write(lowerCAmelCase_     )

     if not save_generations:
        return

     if "preds" in metrics:
        __lowerCAmelCase								=	'\n'.join(metrics['preds']     )
        generations_file.open('w+'     ).write(lowerCAmelCase_     )






  @rank_zero_only
  def 				lowercase		(				self :		Union[str, Any]					,  lowerCAmelCase_ :		Union[str, Any]					,  lowerCAmelCase_ :		List[str]     )							->						Dict:
     try:
        __lowerCAmelCase								=	pl_module.model.model.num_parameters()
     except AttributeError:
        __lowerCAmelCase								=	pl_module.model.num_parameters()

     __lowerCAmelCase								=	count_trainable_parameters(lowerCAmelCase_     )
     # mp stands for million parameters
     trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1e6, 'grad_mp': n_trainable_pars / 1e6}     )






  @rank_zero_only
  def 				lowercase		(				self :		int					,  lowerCAmelCase_ :		pl.Trainer					,  lowerCAmelCase_ :		pl.LightningModule     )							->						Any:
     save_json(pl_module.metrics					,  pl_module.metrics_save_path     )
     return self._write_logs(lowerCAmelCase_					,  lowerCAmelCase_					,  'test'     )






  @rank_zero_only
  def 				lowercase		(				self :		List[Any]					,  lowerCAmelCase_ :		pl.Trainer					,  lowerCAmelCase_ :		Any     )							->						int:
     save_json(pl_module.metrics					,  pl_module.metrics_save_path     )
     # Uncommenting this will save val generations
     # return self._write_logs(trainer, pl_module, "valid")

