



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()
lowercase_	       =  logging.get_logger(__name__)

lowercase_	       =  [
    ('''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'''),
]

lowercase_	       =  [
    '''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		(  __lowerCamelCase	:    Union[str, Any]	)  ->						Dict:
					__lowerCAmelCase	   =torch.load(snake_case__				,   map_location="""cpu"""	)
					return sd






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






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

					__lowerCAmelCase	   =VisualBertConfig(**snake_case__	)

					# Load State Dict
					__lowerCAmelCase	   =load_state_dict(snake_case__	)

					__lowerCAmelCase	   =get_new_dict(snake_case__				,   snake_case__	)

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

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


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









import os
from typing import Dict, List, Tuple, TypeVar, Union


lowercase_	       =			TypeVar('''T''')

lowercase_	       =			Union[List[T], Tuple[T, ...]]
lowercase_	       =			Union[T, List[T], Dict[str, T]]
lowercase_	       =			Union[str, bytes, os.PathLike]




