"""simple docstring"""
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()
SCREAMING_SNAKE_CASE__   =					logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__   =					[
    ("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"),
]

SCREAMING_SNAKE_CASE__   =					[
    "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 				lowerCAmelCase__      (     _UpperCamelCase    :			Any   )   ->			List[str]:





      """simple docstring"""






      snake_case        =     torch.load(lowercase__	,     map_location='cpu'   )
      return sd





def 				lowerCAmelCase__      (     _UpperCamelCase    :			int	,     _UpperCamelCase    :			Optional[Any]	,     _UpperCamelCase    :			List[str]=rename_keys_prefix   )   ->			Any:





      """simple docstring"""






      snake_case        =     OrderedDict()
      snake_case        =     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        =     key
            for name_pair in rename_keys_prefix:
                  snake_case        =     new_key.replace(name_pair[0]	,     name_pair[1]   )
            snake_case        =     d[key]
            if key == "bert.cls.predictions.decoder.weight":
                  # Old bert code didn't have `decoder.bias`, but was added separately
                  snake_case        =     new_d["""cls.predictions.bias"""]
      return new_d





@torch.no_grad()
def 				lowerCAmelCase__      (     _UpperCamelCase    :			Optional[Any]	,     _UpperCamelCase    :			Dict   )   ->			int:





      """simple docstring"""






      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        =     """pretraining"""
            if "vcr" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 5_1_2}
            elif "vqa_advanced" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 2_0_4_8}
            elif "vqa" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 2_0_4_8}
            elif "nlvr" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 1_0_2_4}
            else:
                  raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""   )
      else:
            if "vcr" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 5_1_2}
                  snake_case        =     """multichoice"""
            elif "vqa_advanced" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 2_0_4_8}
                  snake_case        =     """vqa_advanced"""
            elif "vqa" in checkpoint_path:
                  snake_case        =     {"""visual_embedding_dim""": 2_0_4_8, """num_labels""": 3_1_2_9}
                  snake_case        =     """vqa"""
            elif "nlvr" in checkpoint_path:
                  snake_case        =     {
                      """visual_embedding_dim""": 1_0_2_4,
                      """num_labels""": 2,
                  }
                  snake_case        =     """nlvr"""

      snake_case        =     VisualBertConfig(**lowercase__   )

      # Load State Dict
      snake_case        =     load_state_dict(lowercase__   )

      snake_case        =     get_new_dict(lowercase__	,     lowercase__   )

      if model_type == "pretraining":
            snake_case        =     VisualBertForPreTraining(lowercase__   )
      elif model_type == "vqa":
            snake_case        =     VisualBertForQuestionAnswering(lowercase__   )
      elif model_type == "nlvr":
            snake_case        =     VisualBertForVisualReasoning(lowercase__   )
      elif model_type == "multichoice":
            snake_case        =     VisualBertForMultipleChoice(lowercase__   )

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


if __name__ == "__main__":
     SCREAMING_SNAKE_CASE__   =					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.")
     SCREAMING_SNAKE_CASE__   =					parser.parse_args()
     convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)



"""simple docstring"""
import os
import sys
import unittest


SCREAMING_SNAKE_CASE__   =					os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
SCREAMING_SNAKE_CASE__   =					os.path.join(git_repo_path, "src", "transformers")

SCREAMING_SNAKE_CASE__   =					"\n{0} = None\n"

SCREAMING_SNAKE_CASE__   =					"\nclass {0}(metaclass=DummyObject):\n    _backends = {1}\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, {1})\n"


SCREAMING_SNAKE_CASE__   =					"\ndef {0}(*args, **kwargs):\n    requires_backends({0}, {1})\n"



class 					lowerCAmelCase_      (     unittest.TestCase   ):



				"""simple docstring"""

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")'      )
										self.assertIsNone(lowerCAmelCase      )

										snake_case        =     find_backend('    if not is_tokenizers_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tokenizers'      )

										snake_case        =     find_backend('    if not is_tensorflow_text_available():'      )
										self.assertEqual(lowerCAmelCase						,      'tensorflow_text'      )

										snake_case        =     find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tensorflow_text'      )

										snake_case        =     find_backend(
										    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):'      )
										self.assertEqual(lowerCAmelCase						,      'sentencepiece_and_tokenizers_and_vision'      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     read_init()
										# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
										self.assertIn('torch'						,      lowerCAmelCase      )
										self.assertIn('tensorflow_text'						,      lowerCAmelCase      )
										self.assertIn('sentencepiece_and_tokenizers'						,      lowerCAmelCase      )

										# Likewise, we can't assert on the exact content of a key
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertModel'						,      objects['tf']      )
										self.assertIn('FlaxBertModel'						,      objects['flax']      )
										self.assertIn('BertModel'						,      objects['torch']      )
										self.assertIn('TFBertTokenizer'						,      objects['tensorflow_text']      )
										self.assertIn('convert_slow_tokenizer'						,      objects['sentencepiece_and_tokenizers']      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     create_dummy_object('CONSTANT'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      '\nCONSTANT = None\n'      )

										snake_case        =     create_dummy_object('function'						,      '\'torch\''      )
										self.assertEqual(
										    lowerCAmelCase						,      '\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'      )

										snake_case        =     '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
										snake_case        =     create_dummy_object('FakeClass'						,      '\'torch\''      )
										self.assertEqual(lowerCAmelCase						,      lowerCAmelCase      )

				def 						snake_case	(			self      ):






										"""simple docstring"""

										snake_case        =     '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
										snake_case        =     create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']}      )
										self.assertEqual(dummy_files['torch']						,      lowerCAmelCase      )


