


'''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()
__A					:  Any											=      logging.get_logger(__name__)

__A					:  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"),
]

__A					:  Dict											=      [
    "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_			(    A__					:   str					):
  '''simple docstring'''





  lowerCAmelCase_			:      Any			=   torch.load(A__							,    map_location="""cpu"""					)
  return sd







def 				UpperCamelCase_			(    A__					:   Any							,    A__					:   List[str]							,    A__					:   List[Any]=rename_keys_prefix					):
  '''simple docstring'''





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







@torch.no_grad()
def 				UpperCamelCase_			(    A__					:   Union[str, Any]							,    A__					:   Dict					):
  '''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:
    lowerCAmelCase_			:      Union[str, Any]			=   """pretraining"""
    if "vcr" in checkpoint_path:
      lowerCAmelCase_			:      Tuple			=   {"""visual_embedding_dim""": 5_12}
    elif "vqa_advanced" in checkpoint_path:
      lowerCAmelCase_			:      List[Any]			=   {"""visual_embedding_dim""": 20_48}
    elif "vqa" in checkpoint_path:
      lowerCAmelCase_			:      Optional[Any]			=   {"""visual_embedding_dim""": 20_48}
    elif "nlvr" in checkpoint_path:
      lowerCAmelCase_			:      int			=   {"""visual_embedding_dim""": 10_24}
    else:
      raise NotImplementedError(f'No implementation found for `{checkpoint_path}`.'					)
  else:
    if "vcr" in checkpoint_path:
      lowerCAmelCase_			:      Dict			=   {"""visual_embedding_dim""": 5_12}
      lowerCAmelCase_			:      Tuple			=   """multichoice"""
    elif "vqa_advanced" in checkpoint_path:
      lowerCAmelCase_			:      int			=   {"""visual_embedding_dim""": 20_48}
      lowerCAmelCase_			:      List[Any]			=   """vqa_advanced"""
    elif "vqa" in checkpoint_path:
      lowerCAmelCase_			:      Optional[int]			=   {"""visual_embedding_dim""": 20_48, """num_labels""": 31_29}
      lowerCAmelCase_			:      Tuple			=   """vqa"""
    elif "nlvr" in checkpoint_path:
      lowerCAmelCase_			:      Union[str, Any]			=   {
          """visual_embedding_dim""": 10_24,
          """num_labels""": 2,
      }
      lowerCAmelCase_			:      Optional[Any]			=   """nlvr"""

  lowerCAmelCase_			:      str			=   VisualBertConfig(**A__					)

  # Load State Dict
  lowerCAmelCase_			:      Tuple			=   load_state_dict(A__					)

  lowerCAmelCase_			:      Any			=   get_new_dict(A__							,    A__					)

  if model_type == "pretraining":
    lowerCAmelCase_			:      Dict			=   VisualBertForPreTraining(A__					)
  elif model_type == "vqa":
    lowerCAmelCase_			:      List[str]			=   VisualBertForQuestionAnswering(A__					)
  elif model_type == "nlvr":
    lowerCAmelCase_			:      Dict			=   VisualBertForVisualReasoning(A__					)
  elif model_type == "multichoice":
    lowerCAmelCase_			:      Optional[Any]			=   VisualBertForMultipleChoice(A__					)

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


if __name__ == "__main__":
  __A					:  Optional[int]											=      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.")
  __A					:  Union[str, Any]											=      parser.parse_args()
  convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)







'''simple docstring'''






def 				UpperCamelCase_			(    A__					:   str							,    A__					:   list[str]					):
  '''simple docstring'''





  lowerCAmelCase_			:      Dict			=   """"""
  for word_or_phrase in separated:
    if not isinstance(A__							,    A__					):
      raise Exception("""join() accepts only strings to be joined"""					)
    joined += word_or_phrase + separator
  return joined.strip(A__					)


if __name__ == "__main__":
  from doctest import testmod

  testmod()



