






"""simple docstring"""




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 			UpperCamelCase     (      A_ ):


 def _lowercase      (	self						:		int				)   ->			Any:
   _a    :			Tuple			 =      self.config_class(**self.inputs_dict				)
   self.parent.assertTrue(hasattr(snake_case__				,		"""width_multiplier"""				)				)







class 			UpperCamelCase     :


 def __init__(	self						:		str				,		UpperCAmelCase__						:		Tuple				,		UpperCAmelCase__						:		str=13				,		UpperCAmelCase__						:		Dict=64				,		UpperCAmelCase__						:		Any=2				,		UpperCAmelCase__						:		List[Any]=3				,		UpperCAmelCase__						:		Any="swish"				,		UpperCAmelCase__						:		int=3				,		UpperCAmelCase__						:		Union[str, Any]=32				,		UpperCAmelCase__						:		Dict=0.1				,		UpperCAmelCase__						:		List[Any]=0.0_2				,		UpperCAmelCase__						:		Dict=True				,		UpperCAmelCase__						:		Union[str, Any]=True				,		UpperCAmelCase__						:		Dict=10				,		UpperCAmelCase__						:		Optional[Any]=None				,		UpperCAmelCase__						:		List[Any]=0.2_5				,		UpperCAmelCase__						:		int=0.0				,		UpperCAmelCase__						:		List[str]=0.0				,		)   ->			int:
   _a    :			Union[str, Any]			 =      parent
   _a    :			List[str]			 =      batch_size
   _a    :			Dict			 =      image_size
   _a    :			List[Any]			 =      patch_size
   _a    :			List[Any]			 =      num_channels
   _a    :			List[str]			 =      make_divisible(512 * width_multiplier				,		divisor=8				)
   _a    :			Optional[Any]			 =      hidden_act
   _a    :			Optional[Any]			 =      conv_kernel_size
   _a    :			Dict			 =      output_stride
   _a    :			str			 =      classifier_dropout_prob
   _a    :			str			 =      use_labels
   _a    :			int			 =      is_training
   _a    :			List[Any]			 =      num_labels
   _a    :			List[Any]			 =      initializer_range
   _a    :			List[Any]			 =      scope
   _a    :			Optional[int]			 =      width_multiplier
   _a    :			Dict			 =      ffn_dropout
   _a    :			Any			 =      attn_dropout


 def _lowercase      (	self						:		Tuple				)   ->			Optional[int]:
   _a    :			Optional[Any]			 =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

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

   _a    :			Tuple			 =      self.get_config()

   return config, pixel_values, labels, pixel_labels


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


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







@require_torch
class 			UpperCamelCase     (      A_						,	A_						,	unittest.TestCase ):
 UpperCamelCase   :			Optional[int]     =   (
     (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
     if is_torch_available()
     else ()
 )

 UpperCamelCase   :			Optional[int]     =   (
     {
         "feature-extraction": MobileViTVaModel,
         "image-classification": MobileViTVaForImageClassification,
         "image-segmentation": MobileViTVaForSemanticSegmentation,
     }
     if is_torch_available()
     else {}
 )

 UpperCamelCase   :			int     =   False
 UpperCamelCase   :			Union[str, Any]     =   False
 UpperCamelCase   :			int     =   False
 UpperCamelCase   :			str     =   False


 def _lowercase      (	self						:		str				)   ->			Optional[int]:
   _a    :			Tuple			 =      MobileViTVaModelTester(self				)
   _a    :			str			 =      MobileViTVaConfigTester(self				,		config_class=snake_case__				,		has_text_modality=snake_case__				)


 def _lowercase      (	self						:		Any				)   ->			Optional[Any]:
   self.config_tester.run_common_tests()


 @unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""				)
 def _lowercase      (	self						:		Any				)   ->			Dict:
   pass


 @unittest.skip(reason="""MobileViTV2 does not support input and output embeddings"""				)
 def _lowercase      (	self						:		int				)   ->			str:
   pass


 @unittest.skip(reason="""MobileViTV2 does not output attentions"""				)
 def _lowercase      (	self						:		Dict				)   ->			str:
   pass


 @require_torch_multi_gpu
 @unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""				)
 def _lowercase      (	self						:		str				)   ->			List[Any]:
   pass


 @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""				)
 def _lowercase      (	self						:		int				)   ->			Union[str, Any]:
   pass


 def _lowercase      (	self						:		List[Any]				)   ->			Any:
   _a    :			Union[str, Any]			 =      self.model_tester.prepare_config_and_inputs_for_common()

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

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


 def _lowercase      (	self						:		List[str]				)   ->			Optional[int]:
   _a    :			Optional[Any]			 =      self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*snake_case__				)


 def _lowercase      (	self						:		str				)   ->			str:


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

     with torch.no_grad():
       _a    :			List[Any]			 =      model(**self._prepare_for_class(snake_case__				,		snake_case__				)				)

     _a    :			str			 =      outputs.hidden_states

     _a    :			Optional[int]			 =      5
     self.assertEqual(len(snake_case__				)				,		snake_case__				)

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

   _a    :			Dict			 =      self.model_tester.prepare_config_and_inputs_for_common()

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

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

     check_hidden_states_output(snake_case__				,		snake_case__				,		snake_case__				)


 def _lowercase      (	self						:		List[Any]				)   ->			Optional[int]:
   _a    :			Any			 =      self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*snake_case__				)


 def _lowercase      (	self						:		Dict				)   ->			Tuple:
   _a    :			Tuple			 =      self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_semantic_segmentation(*snake_case__				)







 @slow
 def _lowercase      (	self						:		Union[str, Any]				)   ->			Dict:


   for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _a    :			Optional[int]			 =      MobileViTVaModel.from_pretrained(snake_case__				)
     self.assertIsNotNone(snake_case__				)







def      lowerCAmelCase__	(	):



  '''simple docstring'''
  _a    :			int			 =      Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							)
  return image






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


 @cached_property
 def _lowercase      (	self						:		List[str]				)   ->			Dict:
   return (
       MobileViTImageProcessor.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""				)
       if is_vision_available()
       else None
   )


 @slow
 def _lowercase      (	self						:		Dict				)   ->			str:
   _a    :			Optional[Any]			 =      MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""				).to(
       snake_case__				)

   _a    :			Optional[int]			 =      self.default_image_processor
   _a    :			Optional[Any]			 =      prepare_img()
   _a    :			List[str]			 =      image_processor(images=snake_case__				,		return_tensors="""pt"""				).to(snake_case__				)

   # forward pass
   with torch.no_grad():
     _a    :			List[str]			 =      model(**snake_case__				)

   # verify the logits
   _a    :			str			 =      torch.Size((1, 1000)				)
   self.assertEqual(outputs.logits.shape				,		snake_case__				)

   _a    :			Tuple			 =      torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01]				).to(snake_case__				)

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


 @slow
 def _lowercase      (	self						:		Any				)   ->			Dict:
   _a    :			Dict			 =      MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""				)
   _a    :			Any			 =      model.to(snake_case__				)

   _a    :			Tuple			 =      MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""				)

   _a    :			Any			 =      prepare_img()
   _a    :			str			 =      image_processor(images=snake_case__				,		return_tensors="""pt"""				).to(snake_case__				)

   # forward pass
   with torch.no_grad():
     _a    :			Union[str, Any]			 =      model(**snake_case__				)
   _a    :			List[Any]			 =      outputs.logits

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

   _a    :			str			 =      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=snake_case__				,		)

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







 @slow
 def _lowercase      (	self						:		Any				)   ->			List[Any]:
   _a    :			Optional[Any]			 =      MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""				)
   _a    :			str			 =      model.to(snake_case__				)

   _a    :			List[Any]			 =      MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""				)

   _a    :			List[Any]			 =      prepare_img()
   _a    :			Optional[int]			 =      image_processor(images=snake_case__				,		return_tensors="""pt"""				).to(snake_case__				)

   # forward pass
   with torch.no_grad():
     _a    :			List[Any]			 =      model(**snake_case__				)

   _a    :			Optional[int]			 =      outputs.logits.detach().cpu()

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

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



import os
from argparse import ArgumentParser, Namespace

from ..data import SingleSentenceClassificationProcessor as Processor
from ..pipelines import TextClassificationPipeline
from ..utils import is_tf_available, is_torch_available, logging
from . import BaseTransformersCLICommand


if not is_tf_available() and not is_torch_available():
   raise RuntimeError("""At least one of PyTorch or TensorFlow 2.0+ should be installed to use CLI training""")

# TF training parameters
__lowerCamelCase				=			False
__lowerCamelCase				=			False


def      UpperCamelCase   (       __lowerCamelCase	:					Namespace       ):
       return TrainCommand(__lowerCamelCase       )






class     UpperCAmelCase						(		A_					):



   @staticmethod
   def 							_SCREAMING_SNAKE_CASE  (snake_case__						:   ArgumentParser					)						->      int:


          '''simple docstring'''


          snake_case			:  Any     							= parser.add_parser("train"       ,		help="CLI tool to train a model on a task."					)

          train_parser.add_argument(
              "--train_data"       ,		type=snake_case__       ,		required=snake_case__       ,		help="path to train (and optionally evaluation) dataset as a csv with tab separated labels and sentences."       ,		)
          train_parser.add_argument(
              "--column_label"       ,		type=snake_case__       ,		default=0       ,		help="Column of the dataset csv file with example labels."					)
          train_parser.add_argument(
              "--column_text"       ,		type=snake_case__       ,		default=1       ,		help="Column of the dataset csv file with example texts."					)
          train_parser.add_argument(
              "--column_id"       ,		type=snake_case__       ,		default=2       ,		help="Column of the dataset csv file with example ids."					)
          train_parser.add_argument(
              "--skip_first_row"       ,		action="store_true"       ,		help="Skip the first row of the csv file (headers)."					)

          train_parser.add_argument("--validation_data"       ,		type=snake_case__       ,		default=""       ,		help="path to validation dataset."					)
          train_parser.add_argument(
              "--validation_split"       ,		type=snake_case__       ,		default=0.1       ,		help="if validation dataset is not provided, fraction of train dataset to use as validation dataset."       ,		)

          train_parser.add_argument("--output"       ,		type=snake_case__       ,		default="./"       ,		help="path to saved the trained model."					)

          train_parser.add_argument(
              "--task"       ,		type=snake_case__       ,		default="text_classification"       ,		help="Task to train the model on."					)
          train_parser.add_argument(
              "--model"       ,		type=snake_case__       ,		default="bert-base-uncased"       ,		help="Model's name or path to stored model."					)
          train_parser.add_argument("--train_batch_size"       ,		type=snake_case__       ,		default=32       ,		help="Batch size for training."					)
          train_parser.add_argument("--valid_batch_size"       ,		type=snake_case__       ,		default=64       ,		help="Batch size for validation."					)
          train_parser.add_argument("--learning_rate"       ,		type=snake_case__       ,		default=3e-5       ,		help="Learning rate."					)
          train_parser.add_argument("--adam_epsilon"       ,		type=snake_case__       ,		default=1e-08       ,		help="Epsilon for Adam optimizer."					)
          train_parser.set_defaults(func=snake_case__					)



   def __init__(self						:   List[Any]       ,		snake_case__						:   Namespace					)						->      Tuple:


          '''simple docstring'''


          snake_case			:  Any     							= logging.get_logger("transformers-cli/training"					)

          snake_case			:  List[Any]     							= "tf" if is_tf_available() else "torch"

          os.makedirs(args.output       ,		exist_ok=snake_case__					)
          snake_case			:  Any     							= args.output

          snake_case			:  List[Any]     							= args.column_label
          snake_case			:  Tuple     							= args.column_text
          snake_case			:  str     							= args.column_id

          self.logger.info(f"""Loading {args.task} pipeline for {args.model}"""					)
          if args.task == "text_classification":
                 snake_case			:  Optional[int]     							= TextClassificationPipeline.from_pretrained(args.model					)
          elif args.task == "token_classification":
                 raise NotImplementedError
          elif args.task == "question_answering":
                 raise NotImplementedError

          self.logger.info(f"""Loading dataset from {args.train_data}"""					)
          snake_case			:  Tuple     							= Processor.create_from_csv(
              args.train_data       ,		column_label=args.column_label       ,		column_text=args.column_text       ,		column_id=args.column_id       ,		skip_first_row=args.skip_first_row       ,		)
          snake_case			:  Optional[Any]     							= None
          if args.validation_data:
                 self.logger.info(f"""Loading validation dataset from {args.validation_data}"""					)
                 snake_case			:  Dict     							= Processor.create_from_csv(
                     args.validation_data       ,		column_label=args.column_label       ,		column_text=args.column_text       ,		column_id=args.column_id       ,		skip_first_row=args.skip_first_row       ,		)

          snake_case			:  Union[str, Any]     							= args.validation_split
          snake_case			:  Optional[Any]     							= args.train_batch_size
          snake_case			:  List[str]     							= args.valid_batch_size
          snake_case			:  List[Any]     							= args.learning_rate
          snake_case			:  List[Any]     							= args.adam_epsilon



   def 							_SCREAMING_SNAKE_CASE  (self						:   int					)						->      List[str]:


          '''simple docstring'''


          if self.framework == "tf":
                 return self.run_tf()
          return self.run_torch()



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


          '''simple docstring'''


          raise NotImplementedError



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


          '''simple docstring'''


          self.pipeline.fit(
              self.train_dataset       ,		validation_data=self.valid_dataset       ,		validation_split=self.validation_split       ,		learning_rate=self.learning_rate       ,		adam_epsilon=self.adam_epsilon       ,		train_batch_size=self.train_batch_size       ,		valid_batch_size=self.valid_batch_size       ,		)

          # Save trained pipeline
          self.pipeline.save_pretrained(self.output					)





