



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

lowerCAmelCase__      =    [
    ('''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__      =    [
    '''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 	snake_case_  (				A_ :		Any						):


							'''simple docstring'''





							_lowerCamelCase						:    str         =			torch.load(A_,    map_location='''cpu'''						)
							return sd






def 	snake_case_  (				A_ :		Tuple,    A_ :		str,    A_ :		str=rename_keys_prefix						):


							'''simple docstring'''





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






@torch.no_grad()
def 	snake_case_  (				A_ :		Optional[int],    A_ :		Union[str, Any]						):


							'''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:
														_lowerCamelCase						:    Union[str, Any]         =			'''pretraining'''
														if "vcr" in checkpoint_path:
																					_lowerCamelCase						:    List[Any]         =			{'''visual_embedding_dim''': 5_12}
														elif "vqa_advanced" in checkpoint_path:
																					_lowerCamelCase						:    Any         =			{'''visual_embedding_dim''': 20_48}
														elif "vqa" in checkpoint_path:
																					_lowerCamelCase						:    Tuple         =			{'''visual_embedding_dim''': 20_48}
														elif "nlvr" in checkpoint_path:
																					_lowerCamelCase						:    Optional[Any]         =			{'''visual_embedding_dim''': 10_24}
														else:
																					raise NotImplementedError(F'''No implementation found for `{checkpoint_path}`.'''						)
							else:
														if "vcr" in checkpoint_path:
																					_lowerCamelCase						:    Dict         =			{'''visual_embedding_dim''': 5_12}
																					_lowerCamelCase						:    List[str]         =			'''multichoice'''
														elif "vqa_advanced" in checkpoint_path:
																					_lowerCamelCase						:    Dict         =			{'''visual_embedding_dim''': 20_48}
																					_lowerCamelCase						:    Tuple         =			'''vqa_advanced'''
														elif "vqa" in checkpoint_path:
																					_lowerCamelCase						:    str         =			{'''visual_embedding_dim''': 20_48, '''num_labels''': 31_29}
																					_lowerCamelCase						:    int         =			'''vqa'''
														elif "nlvr" in checkpoint_path:
																					_lowerCamelCase						:    Any         =			{
																					    '''visual_embedding_dim''': 10_24,
																					    '''num_labels''': 2,
																					}
																					_lowerCamelCase						:    Optional[Any]         =			'''nlvr'''

							_lowerCamelCase						:    str         =			VisualBertConfig(**A_						)

							# Load State Dict
							_lowerCamelCase						:    Optional[int]         =			load_state_dict(A_						)

							_lowerCamelCase						:    Dict         =			get_new_dict(A_,    A_						)

							if model_type == "pretraining":
														_lowerCamelCase						:    int         =			VisualBertForPreTraining(A_						)
							elif model_type == "vqa":
														_lowerCamelCase						:    Union[str, Any]         =			VisualBertForQuestionAnswering(A_						)
							elif model_type == "nlvr":
														_lowerCamelCase						:    int         =			VisualBertForVisualReasoning(A_						)
							elif model_type == "multichoice":
														_lowerCamelCase						:    Optional[int]         =			VisualBertForMultipleChoice(A_						)

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


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







"""simple docstring"""






import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD






def 	snake_case_  (				A_ :		Optional[int]						):


       '''simple docstring'''





       def wrapper(*A_ :		int,    **A_ :		List[str]						):
              _lowerCamelCase						:    Tuple         =			timeit.default_timer()
              _lowerCamelCase						:    Optional[Any]         =			func(*A_,    **A_						)
              _lowerCamelCase						:    Tuple         =			timeit.default_timer() - starttime
              return delta

       _lowerCamelCase						:    Optional[Any]         =			func.__name__

       return wrapper






def 	snake_case_  (				A_ :		dict,    A_ :		List[str]=1_00,    A_ :		Optional[Any]=None						):


       '''simple docstring'''





       _lowerCamelCase						:    Union[str, Any]         =			[]
       _lowerCamelCase						:    Optional[int]         =			seq_shapes or {}
       for i in range(A_						):
              _lowerCamelCase						:    Any         =			{}
              for col_id, (k, v) in enumerate(features.items()						):
                     if isinstance(A_,    _ArrayXD						):
                            _lowerCamelCase						:    str         =			np.random.rand(*v.shape						).astype(v.dtype						)
                     elif isinstance(A_,    datasets.Value						):
                            if v.dtype == "string":
                                   _lowerCamelCase						:    List[str]         =			'''The small grey turtle was surprisingly fast when challenged.'''
                            else:
                                   _lowerCamelCase						:    int         =			np.random.randint(10,    size=1						).astype(v.dtype						).item()
                     elif isinstance(A_,    datasets.Sequence						):
                            while isinstance(A_,    datasets.Sequence						):
                                   _lowerCamelCase						:    Union[str, Any]         =			v.feature
                            _lowerCamelCase						:    Dict         =			seq_shapes[k]
                            _lowerCamelCase						:    int         =			np.random.rand(*A_						).astype(v.dtype						)
                     _lowerCamelCase						:    Optional[Any]         =			data

              dummy_data.append((i, example)						)

       return dummy_data







def 	snake_case_  (				A_ :		List[str],    A_ :		Optional[int],    A_ :		Optional[int]=1_00,    A_ :		Optional[Any]=None						):


       '''simple docstring'''





       _lowerCamelCase						:    Optional[int]         =			generate_examples(A_,    num_examples=A_,    seq_shapes=A_						)

       with ArrowWriter(features=A_,    path=A_						) as writer:
              for key, record in dummy_data:
                     _lowerCamelCase						:    Optional[Any]         =			features.encode_example(A_						)
                     writer.write(A_						)

              _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			writer.finalize()

       if not num_final_examples == num_examples:
              raise ValueError(
                  F'''Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}.'''						)

       _lowerCamelCase						:    List[Any]         =			datasets.Dataset.from_file(filename=A_,    info=datasets.DatasetInfo(features=A_						)						)

       return dataset


