






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

SCREAMING_SNAKE_CASE__				     =			[
    ('''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__				     =			[
    '''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__		(						lowerCamelCase_       :       Optional[int]				):
			__a						:  str									=     torch.load(lowerCamelCase_   ,  map_location='cpu'				)
			return sd

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

@torch.no_grad()
def        UpperCAmelCase__		(						lowerCamelCase_       :       Any   ,  lowerCamelCase_       :       Any				):
			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						:  Dict									=     'pretraining'
						if "vcr" in checkpoint_path:
									__a						:  int									=     {'visual_embedding_dim': 5_1_2}
						elif "vqa_advanced" in checkpoint_path:
									__a						:  int									=     {'visual_embedding_dim': 2_0_4_8}
						elif "vqa" in checkpoint_path:
									__a						:  Tuple									=     {'visual_embedding_dim': 2_0_4_8}
						elif "nlvr" in checkpoint_path:
									__a						:  List[Any]									=     {'visual_embedding_dim': 1_0_2_4}
						else:
									raise NotImplementedError(f'''No implementation found for `{checkpoint_path}`.'''				)
			else:
						if "vcr" in checkpoint_path:
									__a						:  int									=     {'visual_embedding_dim': 5_1_2}
									__a						:  Any									=     'multichoice'
						elif "vqa_advanced" in checkpoint_path:
									__a						:  Any									=     {'visual_embedding_dim': 2_0_4_8}
									__a						:  List[str]									=     'vqa_advanced'
						elif "vqa" in checkpoint_path:
									__a						:  List[Any]									=     {'visual_embedding_dim': 2_0_4_8, 'num_labels': 3_1_2_9}
									__a						:  List[Any]									=     'vqa'
						elif "nlvr" in checkpoint_path:
									__a						:  Optional[int]									=     {
									    'visual_embedding_dim': 1_0_2_4,
									    'num_labels': 2,
									}
									__a						:  Optional[Any]									=     'nlvr'

			__a						:  str									=     VisualBertConfig(**lowerCamelCase_				)

			# Load State Dict
			__a						:  str									=     load_state_dict(lowerCamelCase_				)

			__a						:  str									=     get_new_dict(lowerCamelCase_   ,  lowerCamelCase_				)

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

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


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






from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase:	List[Any]       			=		logging.get_logger(__name__)

_lowercase:	int       			=		{
    '''alibaba-damo/mgp-str-base''': '''https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json''',
}



class        lowerCamelCase__				(   UpperCAmelCase				):
			UpperCamelCase__    							="mgp-str"





			def __init__(       self	:	int  ,  lowercase__	:	List[Any]=[32, 1_28]  ,  lowercase__	:	str=4  ,  lowercase__	:	List[str]=3  ,  lowercase__	:	Tuple=27  ,  lowercase__	:	Optional[Any]=38  ,  lowercase__	:	Optional[Any]=5_02_57  ,  lowercase__	:	Dict=3_05_22  ,  lowercase__	:	Dict=7_68  ,  lowercase__	:	Tuple=12  ,  lowercase__	:	int=12  ,  lowercase__	:	int=4.0  ,  lowercase__	:	Tuple=True  ,  lowercase__	:	List[Any]=False  ,  lowercase__	:	int=1e-5  ,  lowercase__	:	int=0.0  ,  lowercase__	:	Optional[Any]=0.0  ,  lowercase__	:	List[Any]=0.0  ,  lowercase__	:	Union[str, Any]=False  ,  lowercase__	:	str=0.0_2  ,  **lowercase__	:	Optional[int]  ,  ):
										super().__init__(**lowercase__						)

										_lowerCAmelCase							=							image_size
										_lowerCAmelCase							=							patch_size
										_lowerCAmelCase							=							num_channels
										_lowerCAmelCase							=							max_token_length
										_lowerCAmelCase							=							num_character_labels
										_lowerCAmelCase							=							num_bpe_labels
										_lowerCAmelCase							=							num_wordpiece_labels
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							mlp_ratio
										_lowerCAmelCase							=							distilled
										_lowerCAmelCase							=							layer_norm_eps
										_lowerCAmelCase							=							drop_rate
										_lowerCAmelCase							=							qkv_bias
										_lowerCAmelCase							=							attn_drop_rate
										_lowerCAmelCase							=							drop_path_rate
										_lowerCAmelCase							=							output_aa_attentions
										_lowerCAmelCase							=							initializer_range






