




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 						__lowercase		(unittest.TestCase							):




      """simple docstring"""





      def __init__(				self    ,				A    ,				A=1_3    ,				A=3_0    ,				A=2    ,				A=3    ,				A=True    ,				A=True    ,				A=3_2    ,				A=5    ,				A=4    ,				A=3_7    ,				A="gelu"    ,				A=0.1    ,				A=0.1    ,				A=1_0    ,				A=0.02    ,				)		->    Tuple:
         snake_case				:    str           =	parent
         snake_case				:    Dict           =	batch_size
         snake_case				:    Optional[int]           =	image_size
         snake_case				:    List[str]           =	patch_size
         snake_case				:    str           =	num_channels
         snake_case				:    int           =	is_training
         snake_case				:    str           =	use_labels
         snake_case				:    Tuple           =	hidden_size
         snake_case				:    Tuple           =	num_hidden_layers
         snake_case				:    Optional[int]           =	num_attention_heads
         snake_case				:    int           =	intermediate_size
         snake_case				:    List[Any]           =	hidden_act
         snake_case				:    Union[str, Any]           =	hidden_dropout_prob
         snake_case				:    Dict           =	attention_probs_dropout_prob
         snake_case				:    Optional[int]           =	type_sequence_label_size
         snake_case				:    Dict           =	initializer_range

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

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    int           =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

         snake_case				:    Tuple           =	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=A    ,				initializer_range=self.initializer_range    ,				)

         return config, pixel_values

      def      UpperCAmelCase				(				self    ,				A    ,				A     )		->    Union[str, Any]:
         snake_case				:    Union[str, Any]           =	FlaxViTModel(config=A     )
         snake_case				:    Optional[Any]           =	model(A     )
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
         snake_case				:    Any           =	(self.image_size, self.image_size)
         snake_case				:    Optional[Any]           =	(self.patch_size, self.patch_size)
         snake_case				:    Dict           =	(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    ,				A    ,				A     )		->    Union[str, Any]:
         snake_case				:    Optional[Any]           =	self.type_sequence_label_size
         snake_case				:    Dict           =	FlaxViTForImageClassification(config=A     )
         snake_case				:    Tuple           =	model(A     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.type_sequence_label_size)     )

         # test greyscale images
         snake_case				:    str           =	1
         snake_case				:    Dict           =	FlaxViTForImageClassification(A     )

         snake_case				:    List[str]           =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
         snake_case				:    Any           =	model(A     )

      def      UpperCAmelCase				(				self     )		->    Any:
         snake_case				:    Union[str, Any]           =	self.prepare_config_and_inputs()
         (
             (
             snake_case
         )     , (
             snake_case
         )     , 
         )				:    str           =	config_and_inputs
         snake_case				:    Optional[Any]           =	{"""pixel_values""": pixel_values}
         return config, inputs_dict




@require_flax
class 						__lowercase		(UpperCamelCase__					,		unittest.TestCase							):




      """simple docstring"""





      _snake_case												=						(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()

      def      UpperCAmelCase				(				self     )		->    None:
         snake_case				:    Tuple           =	FlaxViTModelTester(self     )
         snake_case				:    List[str]           =	ConfigTester(self    ,				config_class=A    ,				has_text_modality=A    ,				hidden_size=3_7     )

      def      UpperCAmelCase				(				self     )		->    Dict:
         self.config_tester.run_common_tests()

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    Optional[int]           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*A     )

      def      UpperCAmelCase				(				self     )		->    Optional[Any]:
         snake_case				:    int           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*A     )

      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case     , snake_case				:    int           =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            snake_case				:    Optional[int]           =	model_class(A     )
            snake_case				:    Union[str, Any]           =	inspect.signature(model.__call__     )
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            snake_case				:    str           =	[*signature.parameters.keys()]

            snake_case				:    List[Any]           =	["""pixel_values"""]
            self.assertListEqual(arg_names[:1]    ,				A     )

      def      UpperCAmelCase				(				self     )		->    Optional[Any]:
         snake_case     , snake_case				:    List[Any]           =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__     ):
               snake_case				:    Any           =	self._prepare_for_class(A    ,				A     )
               snake_case				:    Optional[Any]           =	model_class(A     )

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

               with self.subTest("""JIT Enabled"""     ):
                  snake_case				:    int           =	model_jitted(**A     ).to_tuple()

               with self.subTest("""JIT Disabled"""     ):
                  with jax.disable_jit():
                     snake_case				:    str           =	model_jitted(**A     ).to_tuple()

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

      @slow
      def      UpperCAmelCase				(				self     )		->    Any:
         for model_class_name in self.all_model_classes:
            snake_case				:    str           =	model_class_name.from_pretrained("""google/vit-base-patch16-224"""     )
            snake_case				:    int           =	model(np.ones((1, 3, 2_2_4, 2_2_4)     )     )
            self.assertIsNotNone(A     )











from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       List[str]:
   return ConvertCommand(
       args.model_type						,args.tf_checkpoint						,args.pytorch_dump_output						,args.config						,args.finetuning_task_name					)


lowerCamelCase	:	Tuple				=   '\ntransformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires\nTensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.\n'




class 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





      @staticmethod
      def      UpperCAmelCase				(				A     )		->    List[str]:
         snake_case				:    Union[str, Any]           =	parser.add_parser(
             """convert"""    ,				help="""CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints."""    ,				)
         train_parser.add_argument("""--model_type"""    ,				type=A    ,				required=A    ,				help="""Model's type."""     )
         train_parser.add_argument(
             """--tf_checkpoint"""    ,				type=A    ,				required=A    ,				help="""TensorFlow checkpoint path or folder."""     )
         train_parser.add_argument(
             """--pytorch_dump_output"""    ,				type=A    ,				required=A    ,				help="""Path to the PyTorch saved model output."""     )
         train_parser.add_argument("""--config"""    ,				type=A    ,				default=""""""    ,				help="""Configuration file path or folder."""     )
         train_parser.add_argument(
             """--finetuning_task_name"""    ,				type=A    ,				default=A    ,				help="""Optional fine-tuning task name if the TF model was a finetuned model."""    ,				)
         train_parser.set_defaults(func=A     )

      def __init__(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				*A    ,				)		->    List[Any]:
         snake_case				:    Any           =	logging.get_logger("""transformers-cli/converting"""     )

         self._logger.info(f"""Loading model {model_type}"""     )
         snake_case				:    int           =	model_type
         snake_case				:    Any           =	tf_checkpoint
         snake_case				:    int           =	pytorch_dump_output
         snake_case				:    List[str]           =	config
         snake_case				:    Tuple           =	finetuning_task_name

      def      UpperCAmelCase				(				self     )		->    int:
         if self._model_type == "albert":
            try:
               from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "bert":
            try:
               from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "funnel":
            try:
               from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "t5":
            try:
               from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "gpt":
            from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
                convert_openai_checkpoint_to_pytorch,
            )

            convert_openai_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "transfo_xl":
            try:
               from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
                   convert_transfo_xl_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            if "ckpt" in self._tf_checkpoint.lower():
               snake_case				:    List[Any]           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            else:
               snake_case				:    Tuple           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            convert_transfo_xl_checkpoint_to_pytorch(
                A    ,				self._config    ,				self._pytorch_dump_output    ,				A     )
         elif self._model_type == "gpt2":
            try:
               from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
                   convert_gpta_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "xlnet":
            try:
               from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
                   convert_xlnet_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_xlnet_checkpoint_to_pytorch(
                self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output    ,				self._finetuning_task_name     )
         elif self._model_type == "xlm":
            from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
                convert_xlm_checkpoint_to_pytorch,
            )

            convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "lxmert":
            from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
                convert_lxmert_checkpoint_to_pytorch,
            )

            convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "rembert":
            from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
                convert_rembert_tf_checkpoint_to_pytorch,
            )

            convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         else:
            raise ValueError(
                """--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]"""     )





