
"""simple docstring"""




from __future__ import annotations

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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 tensorflow as tf

    from transformers import TFViTForImageClassification, TFViTModel


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor







class     lowercase:



     def __init__(							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE=1_3     ,				__SCREAMING_SNAKE_CASE=3_0     ,				__SCREAMING_SNAKE_CASE=2     ,				__SCREAMING_SNAKE_CASE=3     ,				__SCREAMING_SNAKE_CASE=True     ,				__SCREAMING_SNAKE_CASE=True     ,				__SCREAMING_SNAKE_CASE=3_2     ,				__SCREAMING_SNAKE_CASE=2     ,				__SCREAMING_SNAKE_CASE=4     ,				__SCREAMING_SNAKE_CASE=3_7     ,				__SCREAMING_SNAKE_CASE="gelu"     ,				__SCREAMING_SNAKE_CASE=0.1     ,				__SCREAMING_SNAKE_CASE=0.1     ,				__SCREAMING_SNAKE_CASE=1_0     ,				__SCREAMING_SNAKE_CASE=0.02     ,				__SCREAMING_SNAKE_CASE=3     ,				__SCREAMING_SNAKE_CASE=None     ,				)		->     List[str]:







         """simple docstring"""



         a__		=			parent
         a__		=			batch_size
         a__		=			image_size
         a__		=			patch_size
         a__		=			num_channels
         a__		=			is_training
         a__		=			use_labels
         a__		=			hidden_size
         a__		=			num_hidden_layers
         a__		=			num_attention_heads
         a__		=			intermediate_size
         a__		=			hidden_act
         a__		=			hidden_dropout_prob
         a__		=			attention_probs_dropout_prob
         a__		=			type_sequence_label_size
         a__		=			initializer_range
         a__		=			scope

         # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
         a__		=			(image_size // patch_size) ** 2
         a__		=			num_patches + 1



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







         """simple docstring"""



         a__		=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

         a__		=			None
         if self.use_labels:
             a__		=			ids_tensor([self.batch_size]     ,				self.type_sequence_label_size   )

         a__		=			self.get_config()

         return config, pixel_values, labels



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         return ViTConfig(
             image_size=self.image_size     ,				patch_size=self.patch_size     ,				num_channels=self.num_channels     ,				hidden_size=self.hidden_size     ,				num_hidden_layers=self.num_hidden_layers     ,				num_attention_heads=self.num_attention_heads     ,				intermediate_size=self.intermediate_size     ,				hidden_act=self.hidden_act     ,				hidden_dropout_prob=self.hidden_dropout_prob     ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,				is_decoder=__SCREAMING_SNAKE_CASE     ,				initializer_range=self.initializer_range     ,				)



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )		->     Any:







         """simple docstring"""



         a__		=			TFViTModel(config=__SCREAMING_SNAKE_CASE   )
         a__		=			model(__SCREAMING_SNAKE_CASE     ,				training=__SCREAMING_SNAKE_CASE   )
         self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, self.seq_length, self.hidden_size)   )

         # Test with an image with different size than the one specified in config.
         a__		=			self.image_size // 2
         a__		=			pixel_values[:, :, :image_size, :image_size]
         a__		=			model(__SCREAMING_SNAKE_CASE     ,				interpolate_pos_encoding=__SCREAMING_SNAKE_CASE     ,				training=__SCREAMING_SNAKE_CASE   )
         a__		=			(image_size // self.patch_size) ** 2 + 1
         self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, seq_length, self.hidden_size)   )



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )		->     List[str]:







         """simple docstring"""



         a__		=			self.type_sequence_label_size
         a__		=			TFViTForImageClassification(__SCREAMING_SNAKE_CASE   )
         a__		=			model(__SCREAMING_SNAKE_CASE     ,				labels=__SCREAMING_SNAKE_CASE     ,				training=__SCREAMING_SNAKE_CASE   )
         self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)   )

         # Test with an image with different size than the one specified in config.
         a__		=			self.image_size // 2
         a__		=			pixel_values[:, :, :image_size, :image_size]
         a__		=			model(__SCREAMING_SNAKE_CASE     ,				interpolate_pos_encoding=__SCREAMING_SNAKE_CASE     ,				training=__SCREAMING_SNAKE_CASE   )
         self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)   )

         # test greyscale images
         a__		=			1
         a__		=			TFViTForImageClassification(__SCREAMING_SNAKE_CASE   )
         a__		=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
         a__		=			model(__SCREAMING_SNAKE_CASE   )
         self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)   )




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







         """simple docstring"""



         a__		=			self.prepare_config_and_inputs()
         a__	,     a__	,     a__		=			config_and_inputs
         a__		=			{'pixel_values': pixel_values}
         return config, inputs_dict







@require_tf
class     lowercase(_lowercase    ,						_lowercase    ,						unittest.TestCase       ):
     __snake_case:	Tuple						=     (TFViTModel, TFViTForImageClassification) if is_tf_available() else ()
     __snake_case:	Tuple						=     (
         {'feature-extraction': TFViTModel, 'image-classification': TFViTForImageClassification}
         if is_tf_available()
         else {}
     )

     __snake_case:	Union[str, Any]						=     False
     __snake_case:	Optional[int]						=     False
     __snake_case:	Optional[Any]						=     False



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         a__		=			TFViTModelTester(self   )
         a__		=			ConfigTester(self     ,				config_class=__SCREAMING_SNAKE_CASE     ,				has_text_modality=__SCREAMING_SNAKE_CASE     ,				hidden_size=3_7   )



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







         """simple docstring"""



         self.config_tester.run_common_tests()



     @unittest.skip(reason='ViT does not use inputs_embeds'   )
     def 			lowercase__  (							self   )		->     Optional[int]:







         """simple docstring"""



         pass



     @unittest.skip(reason='ViT does not use inputs_embeds'   )
     def 			lowercase__  (							self   )		->     int:







         """simple docstring"""



         pass



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







         """simple docstring"""



         a__	,     a__		=			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             a__		=			model_class(__SCREAMING_SNAKE_CASE   )
             self.assertIsInstance(model.get_input_embeddings()     ,				(tf.keras.layers.Layer)   )
             a__		=			model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE     ,				tf.keras.layers.Layer   )   )



     def 			lowercase__  (							self   )		->     str:







         """simple docstring"""



         a__	,     a__		=			self.model_tester.prepare_config_and_inputs_for_common()

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

             a__		=			['pixel_values']
             self.assertListEqual(arg_names[:1]     ,				__SCREAMING_SNAKE_CASE   )



     def 			lowercase__  (							self   )		->     int:







         """simple docstring"""



         a__		=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE   )



     def 			lowercase__  (							self   )		->     Any:







         """simple docstring"""



         a__		=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE   )




     @slow
     def 			lowercase__  (							self   )		->     Any:







         """simple docstring"""



         a__		=			TFViTModel.from_pretrained('google/vit-base-patch16-224'   )
         self.assertIsNotNone(__SCREAMING_SNAKE_CASE   )




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







@require_tf
@require_vision
class     lowercase(unittest.TestCase       ):



     @cached_property
     def 			lowercase__  (							self   )		->     Union[str, Any]:







         """simple docstring"""



         return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'   ) if is_vision_available() else None




     @slow
     def 			lowercase__  (							self   )		->     Optional[Any]:







         """simple docstring"""



         a__		=			TFViTForImageClassification.from_pretrained('google/vit-base-patch16-224'   )

         a__		=			self.default_image_processor
         a__		=			prepare_img()
         a__		=			image_processor(images=__SCREAMING_SNAKE_CASE     ,				return_tensors='tf'   )

         # forward pass
         a__		=			model(**__SCREAMING_SNAKE_CASE   )

         # verify the logits
         a__		=			tf.TensorShape((1, 1_0_0_0)   )
         self.assertEqual(outputs.logits.shape     ,				__SCREAMING_SNAKE_CASE   )

         a__		=			tf.constant([-0.27_44, 0.82_15, -0.08_36]   )

         tf.debugging.assert_near(outputs.logits[0, :3]     ,				__SCREAMING_SNAKE_CASE     ,				atol=1e-4   )




"""simple docstring"""




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
a		:       Optional[Any]     						=       False
a		:       str     						=       False



def     __magic_name__		(					UpperCamelCase     :		Namespace     )     ->						Optional[int]:
    return TrainCommand(UpperCamelCase     )







class     lowercase(_lowercase       ):



     @staticmethod
     def 			lowercase__  (							__SCREAMING_SNAKE_CASE   )		->     List[Any]:







         """simple docstring"""



         a__		=			parser.add_parser('train'     ,				help='CLI tool to train a model on a task.'   )

         train_parser.add_argument(
             '--train_data'     ,				type=__SCREAMING_SNAKE_CASE     ,				required=__SCREAMING_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=__SCREAMING_SNAKE_CASE     ,				default=0     ,				help='Column of the dataset csv file with example labels.'   )
         train_parser.add_argument(
             '--column_text'     ,				type=__SCREAMING_SNAKE_CASE     ,				default=1     ,				help='Column of the dataset csv file with example texts.'   )
         train_parser.add_argument(
             '--column_id'     ,				type=__SCREAMING_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=__SCREAMING_SNAKE_CASE     ,				default=''     ,				help='path to validation dataset.'   )
         train_parser.add_argument(
             '--validation_split'     ,				type=__SCREAMING_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=__SCREAMING_SNAKE_CASE     ,				default='./'     ,				help='path to saved the trained model.'   )

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



     def __init__(							self     ,				__SCREAMING_SNAKE_CASE   )		->     Union[str, Any]:







         """simple docstring"""



         a__		=			logging.get_logger('transformers-cli/training'   )

         a__		=			'tf' if is_tf_available() else 'torch'

         os.makedirs(args.output     ,				exist_ok=__SCREAMING_SNAKE_CASE   )
         a__		=			args.output

         a__		=			args.column_label
         a__		=			args.column_text
         a__		=			args.column_id

         self.logger.info(f'Loading {args.task} pipeline for {args.model}'   )
         if args.task == "text_classification":
             a__		=			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}'   )
         a__		=			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     ,				)
         a__		=			None
         if args.validation_data:
             self.logger.info(f'Loading validation dataset from {args.validation_data}'   )
             a__		=			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     ,				)

         a__		=			args.validation_split
         a__		=			args.train_batch_size
         a__		=			args.valid_batch_size
         a__		=			args.learning_rate
         a__		=			args.adam_epsilon



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



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



     def 			lowercase__  (							self   )		->     Any:







         """simple docstring"""



         raise NotImplementedError




     def 			lowercase__  (							self   )		->     Dict:







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


