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 :	List[str]							     =		logging.get_logger(__name__)

SCREAMING_SNAKE_CASE :	Optional[int]							     =		[
    ("""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 :	List[str]							     =		[
    """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   __A		(			_A					):



 """simple docstring"""


 __a       =  torch.load(_A							,				map_location="cpu"					)
 return sd



def   __A		(			_A							,				_A							,				_A=rename_keys_prefix					):



 """simple docstring"""


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



@torch.no_grad()
def   __A		(			_A							,				_A					):



 """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:
  __a       =  "pretraining"
  if "vcr" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 512}
  elif "vqa_advanced" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 2048}
  elif "vqa" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 2048}
  elif "nlvr" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 1024}
  else:
   raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""					)
 else:
  if "vcr" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 512}
   __a       =  "multichoice"
  elif "vqa_advanced" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 2048}
   __a       =  "vqa_advanced"
  elif "vqa" in checkpoint_path:
   __a       =  {"visual_embedding_dim": 2048, "num_labels": 3129}
   __a       =  "vqa"
  elif "nlvr" in checkpoint_path:
   __a       =  {
       "visual_embedding_dim": 1024,
       "num_labels": 2,
   }
   __a       =  "nlvr"

 __a       =  VisualBertConfig(**_A					)

 # Load State Dict
 __a       =  load_state_dict(_A					)

 __a       =  get_new_dict(_A							,				_A					)

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

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


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



import os
import sys
from contextlib import contextmanager


# Windows only
if os.name == "nt":
  import ctypes
  import msvcrt  # noqa
  class    A_ (       ctypes.Structure ):
    # _fields is a specific attr expected by ctypes
    _SCREAMING_SNAKE_CASE    				=			[("""size""", ctypes.c_int), ("""visible""", ctypes.c_byte)]



def   __A		(			):



 """simple docstring"""


 if os.name == "nt":
  __a       =  CursorInfo()
  __a       =  ctypes.windll.kernelaa.GetStdHandle(-11					)
  ctypes.windll.kernelaa.GetConsoleCursorInfo(_A							,				ctypes.byref(_A					)					)
  __a       =  False
  ctypes.windll.kernelaa.SetConsoleCursorInfo(_A							,				ctypes.byref(_A					)					)
 elif os.name == "posix":
  sys.stdout.write("\033[?25l"					)
  sys.stdout.flush()



def   __A		(			):



 """simple docstring"""


 if os.name == "nt":
  __a       =  CursorInfo()
  __a       =  ctypes.windll.kernelaa.GetStdHandle(-11					)
  ctypes.windll.kernelaa.GetConsoleCursorInfo(_A							,				ctypes.byref(_A					)					)
  __a       =  True
  ctypes.windll.kernelaa.SetConsoleCursorInfo(_A							,				ctypes.byref(_A					)					)
 elif os.name == "posix":
  sys.stdout.write("\033[?25h"					)
  sys.stdout.flush()






@contextmanager
def   __A		(			):



 """simple docstring"""


 try:
  hide_cursor()
  yield
 finally:
  show_cursor()


