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

_lowercase :  List[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"),
]

_lowercase :  Tuple    							=       [
    "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_		(				__SCREAMING_SNAKE_CASE					: Union[str, Any]							):



						"""simple docstring"""





						lowercase_	:     Optional[Any]													=       torch.load(__SCREAMING_SNAKE_CASE      ,      map_location='''cpu'''							)
						return sd
def 					snake_case_		(				__SCREAMING_SNAKE_CASE					: Optional[int]      ,      __SCREAMING_SNAKE_CASE					: List[str]      ,      __SCREAMING_SNAKE_CASE					: Optional[int]=rename_keys_prefix							):



						"""simple docstring"""





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



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

						lowercase_	:     Union[str, Any]													=       VisualBertConfig(**__SCREAMING_SNAKE_CASE							)

						# Load State Dict
						lowercase_	:     str													=       load_state_dict(__SCREAMING_SNAKE_CASE							)

						lowercase_	:     int													=       get_new_dict(__SCREAMING_SNAKE_CASE      ,      __SCREAMING_SNAKE_CASE							)

						if model_type == "pretraining":
												lowercase_	:     str													=       VisualBertForPreTraining(__SCREAMING_SNAKE_CASE							)
						elif model_type == "vqa":
												lowercase_	:     Union[str, Any]													=       VisualBertForQuestionAnswering(__SCREAMING_SNAKE_CASE							)
						elif model_type == "nlvr":
												lowercase_	:     List[str]													=       VisualBertForVisualReasoning(__SCREAMING_SNAKE_CASE							)
						elif model_type == "multichoice":
												lowercase_	:     int													=       VisualBertForMultipleChoice(__SCREAMING_SNAKE_CASE							)

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


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







"""simple docstring"""







import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


lowerCAmelCase__      				=   logging.get_logger(__name__)

lowerCAmelCase__      				=   '''T5Config'''
def 					a__			(				SCREAMING_SNAKE_CASE    :    jnp.array				,      SCREAMING_SNAKE_CASE    :    int				,      SCREAMING_SNAKE_CASE    :    int				):


  '''simple docstring'''







  lowerCAmelCase     :     List[str]								=							jnp.zeros_like(SCREAMING_SNAKE_CASE				)
  lowerCAmelCase     :     Optional[int]								=							shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]				)
  lowerCAmelCase     :     List[str]								=							shifted_input_ids.at[:, 0].set(SCREAMING_SNAKE_CASE				)

  lowerCAmelCase     :     str								=							jnp.where(shifted_input_ids == -1_0_0				,      SCREAMING_SNAKE_CASE				,      SCREAMING_SNAKE_CASE				)
  return shifted_input_ids


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				List[Any]   ="mt5"
       a  :				Tuple   =MTaConfig


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				Union[str, Any]   ="mt5"
       a  :				Optional[Any]   =MTaConfig


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				str   ="mt5"
       a  :				Dict   =MTaConfig





