






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

__lowerCamelCase   :					Tuple   							=					[
    ('''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'''),
]

__lowerCamelCase   :					Optional[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       ( __magic_name__						)->    str:



			"""simple docstring"""



			snake_case_				:					List[Any]   =       torch.load(__magic_name__     ,map_location="cpu"						)
			return sd

def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__     ,__magic_name__=rename_keys_prefix						)->    Optional[int]:



			"""simple docstring"""



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

@torch.no_grad()
def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__						)->    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:
						snake_case_				:					List[str]   =       "pretraining"
						if "vcr" in checkpoint_path:
									snake_case_				:					Union[str, Any]   =       {"visual_embedding_dim": 512}
						elif "vqa_advanced" in checkpoint_path:
									snake_case_				:					List[str]   =       {"visual_embedding_dim": 2048}
						elif "vqa" in checkpoint_path:
									snake_case_				:					Tuple   =       {"visual_embedding_dim": 2048}
						elif "nlvr" in checkpoint_path:
									snake_case_				:					List[str]   =       {"visual_embedding_dim": 1024}
						else:
									raise NotImplementedError(F'''No implementation found for `{checkpoint_path}`.'''						)
			else:
						if "vcr" in checkpoint_path:
									snake_case_				:					List[Any]   =       {"visual_embedding_dim": 512}
									snake_case_				:					str   =       "multichoice"
						elif "vqa_advanced" in checkpoint_path:
									snake_case_				:					Tuple   =       {"visual_embedding_dim": 2048}
									snake_case_				:					Union[str, Any]   =       "vqa_advanced"
						elif "vqa" in checkpoint_path:
									snake_case_				:					Optional[int]   =       {"visual_embedding_dim": 2048, "num_labels": 3129}
									snake_case_				:					Union[str, Any]   =       "vqa"
						elif "nlvr" in checkpoint_path:
									snake_case_				:					str   =       {
									    "visual_embedding_dim": 1024,
									    "num_labels": 2,
									}
									snake_case_				:					Dict   =       "nlvr"

			snake_case_				:					List[str]   =       VisualBertConfig(**__magic_name__						)

			# Load State Dict
			snake_case_				:					Any   =       load_state_dict(__magic_name__						)

			snake_case_				:					Dict   =       get_new_dict(__magic_name__     ,__magic_name__						)

			if model_type == "pretraining":
						snake_case_				:					Dict   =       VisualBertForPreTraining(__magic_name__						)
			elif model_type == "vqa":
						snake_case_				:					List[Any]   =       VisualBertForQuestionAnswering(__magic_name__						)
			elif model_type == "nlvr":
						snake_case_				:					Union[str, Any]   =       VisualBertForVisualReasoning(__magic_name__						)
			elif model_type == "multichoice":
						snake_case_				:					List[Any]   =       VisualBertForMultipleChoice(__magic_name__						)

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


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












'''simple docstring'''




import pickle
import unittest

import torch

from accelerate import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils import require_cpu







@require_cpu
class    A_	(unittest.TestCase							):
							"""simple docstring"""




							def 							_A	(		self			:str				)				->  Union[str, Any]:


										'''simple docstring'''




										snake_case_				:					Union[str, Any]   =       torch.nn.Linear(10       ,							10				)
										snake_case_				:					Dict   =       torch.optim.SGD(model.parameters()       ,							0.1				)
										snake_case_				:					Tuple   =       Accelerator()
										snake_case_				:					Optional[Any]   =       accelerator.prepare(lowerCAmelCase__				)
										try:
													pickle.loads(pickle.dumps(lowerCAmelCase__				)				)
										except Exception as e:
													self.fail(F'''Accelerated optimizer pickling failed with {e}'''				)
										AcceleratorState._reset_state()




