






"""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	:	Union[str, Any]				=      logging.get_logger(__name__)

_lowerCAmelCase	:	Optional[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"""),
]

_lowerCAmelCase	:	int				=      [
    """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       : Union[str, Any]				)->		Optional[Any]:






     '''simple docstring'''



     UpperCAmelCase__						:    Any							=							torch.load(snake_case			,					map_location="cpu"				)
     return sd


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






     '''simple docstring'''



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


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

     UpperCAmelCase__						:    Dict							=							VisualBertConfig(**snake_case				)

     # Load State Dict
     UpperCAmelCase__						:    Tuple							=							load_state_dict(snake_case				)

     UpperCAmelCase__						:    Optional[Any]							=							get_new_dict(snake_case			,					snake_case				)

     if model_type == "pretraining":
          UpperCAmelCase__						:    Optional[int]							=							VisualBertForPreTraining(snake_case				)
     elif model_type == "vqa":
          UpperCAmelCase__						:    Optional[Any]							=							VisualBertForQuestionAnswering(snake_case				)
     elif model_type == "nlvr":
          UpperCAmelCase__						:    List[str]							=							VisualBertForVisualReasoning(snake_case				)
     elif model_type == "multichoice":
          UpperCAmelCase__						:    List[Any]							=							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__":
 _lowerCAmelCase	:	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.""")
 _lowerCAmelCase	:	List[Any]				=      parser.parse_args()
 convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)










"""simple docstring"""



import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class 			lowerCAmelCase__			(    __magic_name__							,						__magic_name__							,						__magic_name__   ):
   @register_to_config
   def __init__(		self      :	List[Any]							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	float							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	str							,			snake_case__      :	bool = False							,			):


        '''simple docstring'''
        super().__init__()

        UpperCAmelCase__						:    Optional[int]							=							nn.Embedding(snake_case__							,			snake_case__							)

        UpperCAmelCase__						:    List[Any]							=							nn.Embedding(snake_case__							,			snake_case__							)
        UpperCAmelCase__						:    Optional[Any]							=							False

        UpperCAmelCase__						:    Optional[int]							=							nn.Dropout(p=snake_case__							)

        UpperCAmelCase__						:    Optional[int]							=							TaConfig(
            vocab_size=snake_case__							,			d_model=snake_case__							,			num_heads=snake_case__							,			d_kv=snake_case__							,			d_ff=snake_case__							,			dropout_rate=snake_case__							,			feed_forward_proj=snake_case__							,			is_decoder=snake_case__							,			is_encoder_decoder=snake_case__							,			)

        UpperCAmelCase__						:    Tuple							=							nn.ModuleList()
        for lyr_num in range(snake_case__							):
             UpperCAmelCase__						:    Tuple							=							TaBlock(snake_case__							)
             self.encoders.append(snake_case__							)

        UpperCAmelCase__						:    str							=							TaLayerNorm(snake_case__							)
        UpperCAmelCase__						:    Tuple							=							nn.Dropout(p=snake_case__							)




   def       __a							(		self      :	int							,			snake_case__      :	List[str]							,			snake_case__      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							self.token_embedder(snake_case__							)

        UpperCAmelCase__						:    Optional[int]							=							encoder_input_tokens.shape[1]
        UpperCAmelCase__						:    List[Any]							=							torch.arange(snake_case__							,			device=encoder_input_tokens.device							)
        x += self.position_encoding(snake_case__							)

        UpperCAmelCase__						:    Union[str, Any]							=							self.dropout_pre(snake_case__							)

        # inverted the attention mask
        UpperCAmelCase__						:    List[Any]							=							encoder_input_tokens.size()
        UpperCAmelCase__						:    Tuple							=							self.get_extended_attention_mask(snake_case__							,			snake_case__							)

        for lyr in self.encoders:
             UpperCAmelCase__						:    Any							=							lyr(snake_case__							,			snake_case__							)[0]
        UpperCAmelCase__						:    Any							=							self.layer_norm(snake_case__							)

        return self.dropout_post(snake_case__							), encoder_inputs_mask


