'''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: Dict          =   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: List[Any]          =   [
    """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       (   UpperCAmelCase     ):
 lowercase__				:    Dict  =					torch.load(UpperCAmelCase		,     map_location='''cpu'''     )
 return sd
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=rename_keys_prefix     ):
 lowercase__				:    Union[str, Any]  =					OrderedDict()
 lowercase__				:    Union[str, Any]  =					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
  lowercase__				:    Optional[int]  =					key
  for name_pair in rename_keys_prefix:
   lowercase__				:    List[str]  =					new_key.replace(name_pair[0]		,     name_pair[1]     )
  lowercase__				:    Tuple  =					d[key]
  if key == "bert.cls.predictions.decoder.weight":
   # Old bert code didn't have `decoder.bias`, but was added separately
   lowercase__				:    str  =					new_d['''cls.predictions.bias''']
 return new_d
@torch.no_grad()
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase     ):
 assert (
     checkpoint_path.split('''/'''     )[-1] in ACCEPTABLE_CHECKPOINTS
 ), F"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

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

 lowercase__				:    List[Any]  =					VisualBertConfig(**UpperCAmelCase     )

 # Load State Dict
 lowercase__				:    Optional[Any]  =					load_state_dict(UpperCAmelCase     )

 lowercase__				:    Union[str, Any]  =					get_new_dict(UpperCAmelCase		,     UpperCAmelCase     )

 if model_type == "pretraining":
  lowercase__				:    Optional[Any]  =					VisualBertForPreTraining(UpperCAmelCase     )
 elif model_type == "vqa":
  lowercase__				:    List[Any]  =					VisualBertForQuestionAnswering(UpperCAmelCase     )
 elif model_type == "nlvr":
  lowercase__				:    Optional[Any]  =					VisualBertForVisualReasoning(UpperCAmelCase     )
 elif model_type == "multichoice":
  lowercase__				:    Union[str, Any]  =					VisualBertForMultipleChoice(UpperCAmelCase     )

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


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


'''simple docstring'''

def 			__UpperCamelCase       (   UpperCAmelCase     ):
 lowercase__				:    List[str]  =					0
 while num > 0:
  digit_sum += num % 10
  num //= 10
 return digit_sum
def 			__UpperCamelCase       (   UpperCAmelCase = 100     ):
 lowercase__				:    Dict  =					1
 lowercase__				:    Optional[int]  =					2

 for i in range(2		,     max_n + 1     ):
  lowercase__				:    List[str]  =					pre_numerator
  lowercase__				:    Optional[Any]  =					2 * i // 3 if i % 3 == 0 else 1
  lowercase__				:    Union[str, Any]  =					cur_numerator
  lowercase__				:    int  =					e_cont * pre_numerator + temp

 return sum_digits(UpperCAmelCase     )


if __name__ == "__main__":
  print(F'{solution() = }')

