
"""simple docstring"""

import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
      import jax

      from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel

class __lowerCAmelCase		(				unittest.TestCase      ):





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase=13     ,					__UpperCAmelCase=30     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=3     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=32     ,					__UpperCAmelCase=5     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase=37     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=10     ,					__UpperCAmelCase=0.0_2     ,					):

      '''simple docstring'''




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

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





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




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

      __UpperCamelCase      =	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=__UpperCAmelCase     ,					initializer_range=self.initializer_range     ,					)

      return config, pixel_values





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	FlaxViTModel(config=__UpperCAmelCase				)
      __UpperCamelCase      =	model(__UpperCAmelCase				)
      # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
      __UpperCamelCase      =	(self.image_size, self.image_size)
      __UpperCamelCase      =	(self.patch_size, self.patch_size)
      __UpperCamelCase      =	(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
      self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, num_patches + 1, self.hidden_size)				)





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	self.type_sequence_label_size
      __UpperCamelCase      =	FlaxViTForImageClassification(config=__UpperCAmelCase				)
      __UpperCamelCase      =	model(__UpperCAmelCase				)
      self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, self.type_sequence_label_size)				)

      # test greyscale images
      __UpperCamelCase      =	1
      __UpperCamelCase      =	FlaxViTForImageClassification(__UpperCAmelCase				)

      __UpperCamelCase      =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
      __UpperCamelCase      =	model(__UpperCAmelCase				)







     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




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


@require_flax
class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ,  unittest.TestCase      ):
     lowercase						       =						(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	FlaxViTModelTester(self				)
      __UpperCamelCase      =	ConfigTester(self     ,					config_class=__UpperCAmelCase     ,					has_text_modality=__UpperCAmelCase     ,					hidden_size=37				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs_for_common()

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

       __UpperCamelCase      =	['pixel_values']
       self.assertListEqual(arg_names[:1]     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__				):
        __UpperCamelCase      =	self._prepare_for_class(__UpperCAmelCase     ,					__UpperCAmelCase				)
        __UpperCamelCase      =	model_class(__UpperCAmelCase				)

        @jax.jit
        def model_jitted(__UpperCAmelCase     ,					**__UpperCAmelCase				):
         return model(pixel_values=__UpperCAmelCase     ,					**__UpperCAmelCase				)

        with self.subTest('JIT Enabled'				):
         __UpperCamelCase      =	model_jitted(**__UpperCAmelCase				).to_tuple()

        with self.subTest('JIT Disabled'				):
         with jax.disable_jit():
          __UpperCamelCase      =	model_jitted(**__UpperCAmelCase				).to_tuple()

        self.assertEqual(len(__UpperCAmelCase				)     ,					len(__UpperCAmelCase				)				)
        for jitted_output, output in zip(__UpperCAmelCase     ,					__UpperCAmelCase				):
         self.assertEqual(jitted_output.shape     ,					output.shape				)







     @slow
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      for model_class_name in self.all_model_classes:
       __UpperCamelCase      =	model_class_name.from_pretrained('google/vit-base-patch16-224'				)
       __UpperCamelCase      =	model(np.ones((1, 3, 224, 224)				)				)
       self.assertIsNotNone(__UpperCAmelCase				)





"""simple docstring"""

from argparse import ArgumentParser

from .add_new_model import AddNewModelCommand
from .add_new_model_like import AddNewModelLikeCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .pt_to_tf import PTtoTFCommand
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands






def 							A						(    )				->     Union[str, Any]:
 __UpperCamelCase      =	ArgumentParser('Transformers CLI tool'    , usage='transformers-cli <command> [<args>]'		)
 __UpperCamelCase      =	parser.add_subparsers(help='transformers-cli command helpers'		)

 # Register commands
 ConvertCommand.register_subcommand(snake_case		)
 DownloadCommand.register_subcommand(snake_case		)
 EnvironmentCommand.register_subcommand(snake_case		)
 RunCommand.register_subcommand(snake_case		)
 ServeCommand.register_subcommand(snake_case		)
 UserCommands.register_subcommand(snake_case		)
 AddNewModelCommand.register_subcommand(snake_case		)
 AddNewModelLikeCommand.register_subcommand(snake_case		)
 LfsCommands.register_subcommand(snake_case		)
 PTtoTFCommand.register_subcommand(snake_case		)

 # Let's go
 __UpperCamelCase      =	parser.parse_args()

 if not hasattr(snake_case    , 'func'		):
  parser.print_help()
  exit(1		)

 # Run
 __UpperCamelCase      =	args.func(snake_case		)
 service.run()


if __name__ == "__main__":
      main()



