






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

SCREAMING_SNAKE_CASE							=						[
    ('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'),
]

SCREAMING_SNAKE_CASE							=						[
    '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 				lowercase_					(     __A			:							Union[str, Any]      )							->   List[Any]:



 """simple docstring"""




 lowercase       : str						=torch.load(__A       ,   map_location='''cpu'''      )
 return sd


def 				lowercase_					(     __A			:							Optional[Any]       ,   __A			:							Union[str, Any]       ,   __A			:							Tuple=rename_keys_prefix      )							->   Tuple:



 """simple docstring"""




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


@torch.no_grad()
def 				lowercase_					(     __A			:							Dict       ,   __A			:							Optional[int]      )							->   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:
  lowercase       : Union[str, Any]						='''pretraining'''
  if "vcr" in checkpoint_path:
   lowercase       : Optional[int]						={'''visual_embedding_dim''': 5_1_2}
  elif "vqa_advanced" in checkpoint_path:
   lowercase       : Optional[Any]						={'''visual_embedding_dim''': 2_0_4_8}
  elif "vqa" in checkpoint_path:
   lowercase       : str						={'''visual_embedding_dim''': 2_0_4_8}
  elif "nlvr" in checkpoint_path:
   lowercase       : Tuple						={'''visual_embedding_dim''': 1_0_2_4}
  else:
   raise NotImplementedError(F'No implementation found for `{checkpoint_path}`.'      )
 else:
  if "vcr" in checkpoint_path:
   lowercase       : Optional[int]						={'''visual_embedding_dim''': 5_1_2}
   lowercase       : Dict						='''multichoice'''
  elif "vqa_advanced" in checkpoint_path:
   lowercase       : Optional[int]						={'''visual_embedding_dim''': 2_0_4_8}
   lowercase       : int						='''vqa_advanced'''
  elif "vqa" in checkpoint_path:
   lowercase       : Optional[int]						={'''visual_embedding_dim''': 2_0_4_8, '''num_labels''': 3_1_2_9}
   lowercase       : Any						='''vqa'''
  elif "nlvr" in checkpoint_path:
   lowercase       : Optional[int]						={
       '''visual_embedding_dim''': 1_0_2_4,
       '''num_labels''': 2,
   }
   lowercase       : str						='''nlvr'''

 lowercase       : Dict						=VisualBertConfig(**__A      )

 # Load State Dict
 lowercase       : Union[str, Any]						=load_state_dict(__A      )

 lowercase       : Optional[Any]						=get_new_dict(__A       ,   __A      )

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

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


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











'''simple docstring'''



from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
      from .deepspeed import (
          DeepSpeedEngineWrapper,
          DeepSpeedOptimizerWrapper,
          DeepSpeedSchedulerWrapper,
          DummyOptim,
          DummyScheduler,
          HfDeepSpeedConfig,
      )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers



