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

lowercase  :					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'),
]

lowercase  :					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 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Dict)					->    Optional[Any]:






							'''simple docstring'''





							__UpperCamelCase	:   str			   =					torch.load(_lowerCamelCase						,  map_location="cpu")
							return sd




def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Any						,  _lowerCamelCase :			Optional[int]						,  _lowerCamelCase :			Any=rename_keys_prefix)					->    Union[str, Any]:






							'''simple docstring'''





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




@torch.no_grad()
def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			int						,  _lowerCamelCase :			Optional[Any])					->    List[str]:






							'''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:
														__UpperCamelCase	:   Dict			   =					"pretraining"
														if "vcr" in checkpoint_path:
																					__UpperCamelCase	:   Dict			   =					{"visual_embedding_dim": 512}
														elif "vqa_advanced" in checkpoint_path:
																					__UpperCamelCase	:   List[Any]			   =					{"visual_embedding_dim": 2_048}
														elif "vqa" in checkpoint_path:
																					__UpperCamelCase	:   int			   =					{"visual_embedding_dim": 2_048}
														elif "nlvr" in checkpoint_path:
																					__UpperCamelCase	:   List[Any]			   =					{"visual_embedding_dim": 1_024}
														else:
																					raise NotImplementedError(F'No implementation found for `{checkpoint_path}`.')
							else:
														if "vcr" in checkpoint_path:
																					__UpperCamelCase	:   Any			   =					{"visual_embedding_dim": 512}
																					__UpperCamelCase	:   Dict			   =					"multichoice"
														elif "vqa_advanced" in checkpoint_path:
																					__UpperCamelCase	:   str			   =					{"visual_embedding_dim": 2_048}
																					__UpperCamelCase	:   Tuple			   =					"vqa_advanced"
														elif "vqa" in checkpoint_path:
																					__UpperCamelCase	:   Optional[Any]			   =					{"visual_embedding_dim": 2_048, "num_labels": 3_129}
																					__UpperCamelCase	:   str			   =					"vqa"
														elif "nlvr" in checkpoint_path:
																					__UpperCamelCase	:   Any			   =					{
																					    "visual_embedding_dim": 1_024,
																					    "num_labels": 2,
																					}
																					__UpperCamelCase	:   Dict			   =					"nlvr"

							__UpperCamelCase	:   Dict			   =					VisualBertConfig(**_lowerCamelCase)

							# Load State Dict
							__UpperCamelCase	:   List[str]			   =					load_state_dict(_lowerCamelCase)

							__UpperCamelCase	:   Tuple			   =					get_new_dict(_lowerCamelCase						,  _lowerCamelCase)

							if model_type == "pretraining":
														__UpperCamelCase	:   Optional[int]			   =					VisualBertForPreTraining(_lowerCamelCase)
							elif model_type == "vqa":
														__UpperCamelCase	:   Optional[Any]			   =					VisualBertForQuestionAnswering(_lowerCamelCase)
							elif model_type == "nlvr":
														__UpperCamelCase	:   Any			   =					VisualBertForVisualReasoning(_lowerCamelCase)
							elif model_type == "multichoice":
														__UpperCamelCase	:   Tuple			   =					VisualBertForMultipleChoice(_lowerCamelCase)

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


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

def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			int						,  _lowerCamelCase :			int)					->    float:






							'''simple docstring'''





							return base * power(_lowerCamelCase						,  (exponent - 1)) if exponent else 1


if __name__ == "__main__":
		print('Raise base to the power of exponent using recursion...')
		lowercase  :					Optional[int]													=       int(input('Enter the base: ').strip())
		lowercase  :					Tuple													=       int(input('Enter the exponent: ').strip())
		lowercase  :					str													=       power(base, abs(exponent))
		if exponent < 0:  # power() does not properly deal w/ negative exponents
				lowercase  :					Tuple													=       1 / result
		print(f"{base} to the power of {exponent} is {result}")