
from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
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 TFResNetForImageClassification, TFResNetModel
   from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor






class     UpperCAmelCase						:



   def __init__(self						:   List[Any]       ,		snake_case__						:   str       ,		snake_case__						:   int=3       ,		snake_case__						:   Any=32       ,		snake_case__						:   Dict=3       ,		snake_case__						:   int=10       ,		snake_case__						:   Optional[Any]=[10, 20, 30, 40]       ,		snake_case__						:   str=[1, 1, 2, 1]       ,		snake_case__						:   str=True       ,		snake_case__						:   Dict=True       ,		snake_case__						:   int="relu"       ,		snake_case__						:   Any=3       ,		snake_case__						:   str=None       ,		)						->      List[Any]:


          '''simple docstring'''


          snake_case			:  int     							= parent
          snake_case			:  Any     							= batch_size
          snake_case			:  Optional[Any]     							= image_size
          snake_case			:  List[Any]     							= num_channels
          snake_case			:  Optional[int]     							= embeddings_size
          snake_case			:  Tuple     							= hidden_sizes
          snake_case			:  str     							= depths
          snake_case			:  int     							= is_training
          snake_case			:  List[str]     							= use_labels
          snake_case			:  Union[str, Any]     							= hidden_act
          snake_case			:  Optional[Any]     							= num_labels
          snake_case			:  str     							= scope
          snake_case			:  str     							= len(snake_case__					)



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


          '''simple docstring'''


          snake_case			:  List[str]     							= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

          snake_case			:  Union[str, Any]     							= None
          if self.use_labels:
                 snake_case			:  Any     							= ids_tensor([self.batch_size]       ,		self.num_labels					)

          snake_case			:  Optional[Any]     							= self.get_config()

          return config, pixel_values, labels



   def 							_SCREAMING_SNAKE_CASE  (self						:   Dict					)						->      List[Any]:


          '''simple docstring'''


          return ResNetConfig(
              num_channels=self.num_channels       ,		embeddings_size=self.embeddings_size       ,		hidden_sizes=self.hidden_sizes       ,		depths=self.depths       ,		hidden_act=self.hidden_act       ,		num_labels=self.num_labels       ,		image_size=self.image_size       ,		)



   def 							_SCREAMING_SNAKE_CASE  (self						:   List[Any]       ,		snake_case__						:   int       ,		snake_case__						:   Tuple       ,		snake_case__						:   int					)						->      Dict:


          '''simple docstring'''


          snake_case			:  List[str]     							= TFResNetModel(config=snake_case__					)
          snake_case			:  List[str]     							= model(snake_case__					)
          # expected last hidden states: B, C, H // 32, W // 32
          self.parent.assertEqual(
              result.last_hidden_state.shape       ,		(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,		)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Optional[Any]       ,		snake_case__						:   List[str]       ,		snake_case__						:   Tuple       ,		snake_case__						:   Any					)						->      Tuple:


          '''simple docstring'''


          snake_case			:  Tuple     							= self.num_labels
          snake_case			:  int     							= TFResNetForImageClassification(snake_case__					)
          snake_case			:  Optional[int]     							= model(snake_case__       ,		labels=snake_case__					)
          self.parent.assertEqual(result.logits.shape       ,		(self.batch_size, self.num_labels)					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   Dict					)						->      List[Any]:


          '''simple docstring'''


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






@require_tf
class     UpperCAmelCase						(		A_			,A_			,unittest.TestCase					):
   A__   :						List[Any]			=	(TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
   A__   :						Dict			=	(
       {"feature-extraction": TFResNetModel, "image-classification": TFResNetForImageClassification}
       if is_tf_available()
       else {}
   )

   A__   :						List[Any]			=	False
   A__   :						Any			=	False
   A__   :						List[str]			=	False
   A__   :						Optional[int]			=	False
   A__   :						Optional[Any]			=	False



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


          '''simple docstring'''


          snake_case			:  Union[str, Any]     							= TFResNetModelTester(self					)
          snake_case			:  Dict     							= ConfigTester(self       ,		config_class=snake_case__       ,		has_text_modality=snake_case__					)



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


          '''simple docstring'''


          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()



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


          '''simple docstring'''


          return



   @unittest.skip(reason="ResNet does not use inputs_embeds"					)
   def 							_SCREAMING_SNAKE_CASE  (self						:   str					)						->      Optional[int]:


          '''simple docstring'''


          pass



   @unittest.skip(reason="ResNet does not support input and output embeddings"					)
   def 							_SCREAMING_SNAKE_CASE  (self						:   int					)						->      Dict:


          '''simple docstring'''


          pass



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


          '''simple docstring'''


          snake_case							,   snake_case			:  Optional[int]     							= self.model_tester.prepare_config_and_inputs_for_common()

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

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



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


          '''simple docstring'''


          snake_case			:  Optional[Any]     							= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case__					)



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


          '''simple docstring'''


          def check_hidden_states_output(snake_case__						:   int       ,		snake_case__						:   List[str]       ,		snake_case__						:   Optional[Any]					):
                 snake_case			:  Optional[int]     							= model_class(snake_case__					)
                 snake_case			:  Any     							= model(**self._prepare_for_class(snake_case__       ,		snake_case__					)					)

                 snake_case			:  Union[str, Any]     							= outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 snake_case			:  List[Any]     							= self.model_tester.num_stages
                 self.assertEqual(len(snake_case__					)       ,		expected_num_stages + 1					)

                 # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
                 self.assertListEqual(
                     list(hidden_states[0].shape[-2:]					)       ,		[self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,		)

          snake_case							,   snake_case			:  Tuple     							= self.model_tester.prepare_config_and_inputs_for_common()
          snake_case			:  Optional[Any]     							= ["basic", "bottleneck"]
          for model_class in self.all_model_classes:
                 for layer_type in layers_type:
                        snake_case			:  List[Any]     							= layer_type
                        snake_case			:  List[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"]
                        snake_case			:  List[str]     							= True

                        check_hidden_states_output(snake_case__       ,		snake_case__       ,		snake_case__					)



   def 							_SCREAMING_SNAKE_CASE  (self						:   str					)						->      Any:


          '''simple docstring'''


          snake_case			:  str     							= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*snake_case__					)



   @slow
   def 							_SCREAMING_SNAKE_CASE  (self						:   List[str]					)						->      Any:


          '''simple docstring'''


          for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 snake_case			:  str     							= TFResNetModel.from_pretrained(snake_case__					)
                 self.assertIsNotNone(snake_case__					)



def      UpperCamelCase   (       ):
       snake_case			:  List[Any]     							= Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
       return image






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



   @cached_property
   def 							_SCREAMING_SNAKE_CASE  (self						:   str					)						->      Optional[int]:


          '''simple docstring'''


          return (
              AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]					)
              if is_vision_available()
              else None
          )



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


          '''simple docstring'''


          snake_case			:  Tuple     							= TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]					)

          snake_case			:  str     							= self.default_image_processor
          snake_case			:  Dict     							= prepare_img()
          snake_case			:  List[str]     							= image_processor(images=snake_case__       ,		return_tensors="tf"					)

          # forward pass
          snake_case			:  List[str]     							= model(**snake_case__					)

          # verify the logits
          snake_case			:  Any     							= tf.TensorShape((1, 10_00)					)
          self.assertEqual(outputs.logits.shape       ,		snake_case__					)

          snake_case			:  str     							= tf.constant([-11.1069, -9.7877, -8.3777]					)

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







import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCamelCase				=			logging.get_logger(__name__)

__lowerCamelCase				=			[
    ("""bert.bert""", """visual_bert"""),
    ("""bert.cls""", """cls"""),
    ("""bert.classifier""", """cls"""),
    ("""token_type_embeddings_visual""", """visual_token_type_embeddings"""),
    ("""position_embeddings_visual""", """visual_position_embeddings"""),
    ("""projection""", """visual_projection"""),
]

__lowerCamelCase				=			[
    """nlvr2_coco_pre_trained.th""",
    """nlvr2_fine_tuned.th""",
    """nlvr2_pre_trained.th""",
    """vcr_coco_pre_train.th""",
    """vcr_fine_tune.th""",
    """vcr_pre_train.th""",
    """vqa_coco_pre_trained.th""",
    """vqa_fine_tuned.th""",
    """vqa_pre_trained.th""",
]


def      UpperCamelCase   (       __lowerCamelCase	:					int       ):
       snake_case			:  List[Any]     							= torch.load(__lowerCamelCase							,   map_location="cpu"       )
       return sd


def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[int]							,   __lowerCamelCase	:					Optional[Any]=rename_keys_prefix       ):
       snake_case			:  Optional[Any]     							= OrderedDict()
       snake_case			:  Optional[int]     							= torch.arange(config.max_position_embeddings       ).expand((1, -1)       )
       # detector_d = OrderedDict()
       for key in d:
              if "detector" in key:
                     # detector_d[key.replace('detector.','')] = d[key]
                     continue
              snake_case			:  Any     							= key
              for name_pair in rename_keys_prefix:
                     snake_case			:  List[Any]     							= new_key.replace(name_pair[0]							,   name_pair[1]       )
              snake_case			:  List[str]     							= d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     snake_case			:  int     							= new_d["cls.predictions.bias"]
       return new_d


@torch.no_grad()
def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					Optional[Any]       ):
       assert (
           checkpoint_path.split("/"       )[-1] in ACCEPTABLE_CHECKPOINTS
       ), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

       # Get Config
       if "pre" in checkpoint_path:
              snake_case			:  Optional[Any]     							= "pretraining"
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[int]     							= {"visual_embedding_dim": 512}
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  int     							= {"visual_embedding_dim": 2048}
              elif "vqa" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
              elif "nlvr" in checkpoint_path:
                     snake_case			:  Tuple     							= {"visual_embedding_dim": 1024}
              else:
                     raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""       )
       else:
              if "vcr" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 512}
                     snake_case			:  List[str]     							= "multichoice"
              elif "vqa_advanced" in checkpoint_path:
                     snake_case			:  Dict     							= {"visual_embedding_dim": 2048}
                     snake_case			:  Tuple     							= "vqa_advanced"
              elif "vqa" in checkpoint_path:
                     snake_case			:  Optional[Any]     							= {"visual_embedding_dim": 2048, "num_labels": 3129}
                     snake_case			:  Tuple     							= "vqa"
              elif "nlvr" in checkpoint_path:
                     snake_case			:  str     							= {
                         "visual_embedding_dim": 1024,
                         "num_labels": 2,
                     }
                     snake_case			:  str     							= "nlvr"

       snake_case			:  int     							= VisualBertConfig(**__lowerCamelCase       )

       # Load State Dict
       snake_case			:  Any     							= load_state_dict(__lowerCamelCase       )

       snake_case			:  List[Any]     							= get_new_dict(__lowerCamelCase							,   __lowerCamelCase       )

       if model_type == "pretraining":
              snake_case			:  str     							= VisualBertForPreTraining(__lowerCamelCase       )
       elif model_type == "vqa":
              snake_case			:  Any     							= VisualBertForQuestionAnswering(__lowerCamelCase       )
       elif model_type == "nlvr":
              snake_case			:  Optional[Any]     							= VisualBertForVisualReasoning(__lowerCamelCase       )
       elif model_type == "multichoice":
              snake_case			:  Optional[int]     							= VisualBertForMultipleChoice(__lowerCamelCase       )

       model.load_state_dict(__lowerCamelCase       )
       # Save Checkpoints
       Path(__lowerCamelCase       ).mkdir(exist_ok=__lowerCamelCase       )
       model.save_pretrained(__lowerCamelCase       )


if __name__ == "__main__":
   __lowerCamelCase				=			argparse.ArgumentParser()
   # Required parameters
   parser.add_argument("""orig_checkpoint_path""", type=str, help="""A path to .th on local filesystem.""")
   parser.add_argument("""pytorch_dump_folder_path""", type=str, help="""Path to the output PyTorch model.""")
   __lowerCamelCase				=			parser.parse_args()
   convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)





