



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

__lowerCamelCase         =			[
    ('''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         =			[
    '''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 (       __snake_case    )     ->     Optional[Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Any				      =		torch.load(__snake_case					,  map_location="cpu"    )
      return sd







def    _snake_case (       __snake_case					,  __snake_case					,  __snake_case=rename_keys_prefix    )     ->     List[Any]:






      '''simple docstring'''







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







@torch.no_grad()
def    _snake_case (       __snake_case					,  __snake_case    )     ->     int:






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

      UpperCAmelCase_  :	Tuple				      =		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[Any]				      =		VisualBertForPreTraining(__snake_case    )
      elif model_type == "vqa":
            UpperCAmelCase_  :	int				      =		VisualBertForQuestionAnswering(__snake_case    )
      elif model_type == "nlvr":
            UpperCAmelCase_  :	Union[str, Any]				      =		VisualBertForVisualReasoning(__snake_case    )
      elif model_type == "multichoice":
            UpperCAmelCase_  :	Union[str, 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         =			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         =			parser.parse_args()
   convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)






from __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class   snake_case_ :




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42


class   snake_case_ :




       """simple docstring"""






       def __init__(     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	list[list[Edge]]				      =		[[] for _ in range(lowercase)]
             UpperCAmelCase_  :	str				      =		size



       def __getitem__(     self     ,lowercase):

             """simple docstring"""





             return iter(self._graph[vertex])



       @property
       def  A_   (     self):

             """simple docstring"""





             return self._size



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             if weight not in (0, 1):
                   raise ValueError("Edge weight must be either 0 or 1.")

             if to_vertex < 0 or to_vertex >= self.size:
                   raise ValueError("Vertex indexes must be in [0; size).")

             self._graph[from_vertex].append(Edge(lowercase     ,lowercase))



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		deque([start_vertex])
             UpperCAmelCase_  :	list[int | None]				      =		[None] * self.size
             UpperCAmelCase_  :	Dict				      =		0

             while queue:
                   UpperCAmelCase_  :	List[Any]				      =		queue.popleft()
                   UpperCAmelCase_  :	Any				      =		distances[current_vertex]
                   if current_distance is None:
                         continue

                   for edge in self[current_vertex]:
                         UpperCAmelCase_  :	Tuple				      =		current_distance + edge.weight
                         UpperCAmelCase_  :	Optional[Any]				      =		distances[edge.destination_vertex]
                         if (
                             isinstance(lowercase     ,lowercase)
                             and new_distance >= dest_vertex_distance
                         ):
                               continue
                         UpperCAmelCase_  :	Tuple				      =		new_distance
                         if edge.weight == 0:
                               queue.appendleft(edge.destination_vertex)
                         else:
                               queue.append(edge.destination_vertex)

             if distances[finish_vertex] is None:
                   raise ValueError("No path from start_vertex to finish_vertex.")

             return distances[finish_vertex]


if __name__ == "__main__":
   import doctest

   doctest.testmod()
