





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

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

UpperCAmelCase					:		Dict		=							[
    '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    _a    (							lowerCAmelCase_	):





      """simple docstring"""



      _snake_case		:	Union[str, Any]					     =     torch.load(lowerCAmelCase_			,      map_location='''cpu'''	)
      return sd




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_=rename_keys_prefix	):





      """simple docstring"""



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




@torch.no_grad()
def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_	):





      """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:
            _snake_case		:	str					     =     '''pretraining'''
            if "vcr" in checkpoint_path:
                  _snake_case		:	List[str]					     =     {'''visual_embedding_dim''': 512}
            elif "vqa_advanced" in checkpoint_path:
                  _snake_case		:	Optional[int]					     =     {'''visual_embedding_dim''': 2_048}
            elif "vqa" in checkpoint_path:
                  _snake_case		:	Tuple					     =     {'''visual_embedding_dim''': 2_048}
            elif "nlvr" in checkpoint_path:
                  _snake_case		:	Dict					     =     {'''visual_embedding_dim''': 1_024}
            else:
                  raise NotImplementedError(f'''No implementation found for `{checkpoint_path}`.'''	)
      else:
            if "vcr" in checkpoint_path:
                  _snake_case		:	str					     =     {'''visual_embedding_dim''': 512}
                  _snake_case		:	Optional[int]					     =     '''multichoice'''
            elif "vqa_advanced" in checkpoint_path:
                  _snake_case		:	int					     =     {'''visual_embedding_dim''': 2_048}
                  _snake_case		:	Dict					     =     '''vqa_advanced'''
            elif "vqa" in checkpoint_path:
                  _snake_case		:	Optional[int]					     =     {'''visual_embedding_dim''': 2_048, '''num_labels''': 3_129}
                  _snake_case		:	Dict					     =     '''vqa'''
            elif "nlvr" in checkpoint_path:
                  _snake_case		:	int					     =     {
                      '''visual_embedding_dim''': 1_024,
                      '''num_labels''': 2,
                  }
                  _snake_case		:	Optional[Any]					     =     '''nlvr'''

      _snake_case		:	str					     =     VisualBertConfig(**lowerCAmelCase_	)

      # Load State Dict
      _snake_case		:	Union[str, Any]					     =     load_state_dict(lowerCAmelCase_	)

      _snake_case		:	Tuple					     =     get_new_dict(lowerCAmelCase_			,      lowerCAmelCase_	)

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

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


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









'''simple docstring'''

import pickle
import unittest

import torch

from accelerate import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils import require_cpu





@require_cpu
class lowerCamelCase						(unittest.TestCase					):







  def 				UpperCAmelCase_				(		self		)		->						Union[str, Any]:






        """simple docstring"""





        _snake_case		:	Any					     =     torch.nn.Linear(10     ,     10		)
        _snake_case		:	Optional[int]					     =     torch.optim.SGD(model.parameters()     ,     0.1		)
        _snake_case		:	List[str]					     =     Accelerator()
        _snake_case		:	Optional[Any]					     =     accelerator.prepare(lowercase__		)
        try:
              pickle.loads(pickle.dumps(lowercase__		)		)
        except Exception as e:
              self.fail(F'''Accelerated optimizer pickling failed with {e}'''		)
        AcceleratorState._reset_state()


