






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()
__a:     Optional[int]   					=					logging.get_logger(__name__)

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

__a:     Union[str, 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 (							__snake_case )				->						Tuple:
       _UpperCAmelCase							=    torch.load(__snake_case   ,      map_location="""cpu""" )
       return sd




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




@torch.no_grad()
def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						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:
              _UpperCAmelCase							=    """pretraining"""
              if "vcr" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 5_1_2}
              elif "vqa_advanced" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 2_0_4_8}
              elif "vqa" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 2_0_4_8}
              elif "nlvr" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 1_0_2_4}
              else:
                     raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`.""" )
       else:
              if "vcr" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 5_1_2}
                     _UpperCAmelCase							=    """multichoice"""
              elif "vqa_advanced" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 2_0_4_8}
                     _UpperCAmelCase							=    """vqa_advanced"""
              elif "vqa" in checkpoint_path:
                     _UpperCAmelCase							=    {"""visual_embedding_dim""": 2_0_4_8, """num_labels""": 3_1_2_9}
                     _UpperCAmelCase							=    """vqa"""
              elif "nlvr" in checkpoint_path:
                     _UpperCAmelCase							=    {
                         """visual_embedding_dim""": 1_0_2_4,
                         """num_labels""": 2,
                     }
                     _UpperCAmelCase							=    """nlvr"""

       _UpperCAmelCase							=    VisualBertConfig(**__snake_case )

       # Load State Dict
       _UpperCAmelCase							=    load_state_dict(__snake_case )

       _UpperCAmelCase							=    get_new_dict(__snake_case   ,      __snake_case )

       if model_type == "pretraining":
              _UpperCAmelCase							=    VisualBertForPreTraining(__snake_case )
       elif model_type == "vqa":
              _UpperCAmelCase							=    VisualBertForQuestionAnswering(__snake_case )
       elif model_type == "nlvr":
              _UpperCAmelCase							=    VisualBertForVisualReasoning(__snake_case )
       elif model_type == "multichoice":
              _UpperCAmelCase							=    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__":
  __a:     List[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.''')
  __a:     Tuple   					=					parser.parse_args()
  convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


a_    :  Tuple						  =		{'configuration_unispeech': ['UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechConfig']}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_    :  Union[str, Any]						  =		[
         'UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST',
         'UniSpeechForCTC',
         'UniSpeechForPreTraining',
         'UniSpeechForSequenceClassification',
         'UniSpeechModel',
         'UniSpeechPreTrainedModel',
     ]

if TYPE_CHECKING:
     from .configuration_unispeech import UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_unispeech import (
              UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST,
              UniSpeechForCTC,
              UniSpeechForPreTraining,
              UniSpeechForSequenceClassification,
              UniSpeechModel,
              UniSpeechPreTrainedModel,
          )

else:
     import sys

     a_    :  Optional[Any]						  =		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)