






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 	UpperCamelCase	(							snake_case__  : Any	):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    torch.load(snake_case__						,map_location="""cpu"""	)
     return sd






def 	UpperCamelCase	(							snake_case__  : List[Any]						,snake_case__  : List[Any]						,snake_case__  : Optional[int]=rename_keys_prefix	):




     '''simple docstring'''






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






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

     __snake_case      :int 							=    VisualBertConfig(**snake_case__	)

     # Load State Dict
     __snake_case      :Union[str, Any] 							=    load_state_dict(snake_case__	)

     __snake_case      :Any 							=    get_new_dict(snake_case__						,snake_case__	)

     if model_type == "pretraining":
          __snake_case      :List[Any] 							=    VisualBertForPreTraining(snake_case__	)
     elif model_type == "vqa":
          __snake_case      :Dict 							=    VisualBertForQuestionAnswering(snake_case__	)
     elif model_type == "nlvr":
          __snake_case      :Optional[int] 							=    VisualBertForVisualReasoning(snake_case__	)
     elif model_type == "multichoice":
          __snake_case      :List[str] 							=    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 collections.abc import Callable




class       snake_case__ :

       '''simple docstring'''



       def __init__(				self      ,    a__ = None     )			->   None:






            '''simple docstring'''





            __snake_case      :list 							=    []
            # Stores indexes of each item for supporting updates and deletion.
            __snake_case      :dict 							=    {}
            # Stores current size of heap.
            __snake_case      :List[Any] 							=    0
            # Stores function used to evaluate the score of an item on which basis ordering
            # will be done.
            __snake_case      :Optional[Any] 							=    key or (lambda a__     : x)



       def        __lowercase    (				self      ,    a__     )			->   int | None:






            '''simple docstring'''





            return int((i - 1) / 2     ) if i > 0 else None



       def        __lowercase    (				self      ,    a__     )			->   int | None:






            '''simple docstring'''





            __snake_case      :int 							=    int(2 * i + 1     )
            return left if 0 < left < self.size else None



       def        __lowercase    (				self      ,    a__     )			->   int | None:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    int(2 * i + 2     )
            return right if 0 < right < self.size else None



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   None:






            '''simple docstring'''





            __snake_case							,			__snake_case      :Dict 							=    (
                self.pos_map[self.arr[j][0]],
                self.pos_map[self.arr[i][0]],
            )
            # Then swap the items in the list.
            __snake_case							,			__snake_case      :Tuple 							=    self.arr[j], self.arr[i]



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   bool:






            '''simple docstring'''





            return self.arr[i][1] < self.arr[j][1]



       def        __lowercase    (				self      ,    a__     )			->   int:






            '''simple docstring'''





            __snake_case      :int 							=    self._left(a__     )
            __snake_case      :Optional[int] 							=    self._right(a__     )
            __snake_case      :List[Any] 							=    i

            if left is not None and not self._cmp(a__      ,    a__     ):
                 __snake_case      :int 							=    left
            if right is not None and not self._cmp(a__      ,    a__     ):
                 __snake_case      :Tuple 							=    right

            return valid_parent



       def        __lowercase    (				self      ,    a__     )			->   None:






            '''simple docstring'''





            __snake_case      :int 							=    self._parent(a__     )
            while parent is not None and not self._cmp(a__      ,    a__     ):
                 self._swap(a__      ,    a__     )
                 __snake_case							,			__snake_case      :str 							=    parent, self._parent(a__     )



       def        __lowercase    (				self      ,    a__     )			->   None:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self._get_valid_parent(a__     )
            while valid_parent != index:
                 self._swap(a__      ,    a__     )
                 __snake_case							,			__snake_case      :Optional[int] 							=    valid_parent, self._get_valid_parent(a__     )



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   None:






            '''simple docstring'''





            if item not in self.pos_map:
                 return
            __snake_case      :List[Any] 							=    self.pos_map[item]
            __snake_case      :Optional[Any] 							=    [item, self.key(a__     )]
            # Make sure heap is right in both up and down direction.
            # Ideally only one of them will make any change.
            self._heapify_up(a__     )
            self._heapify_down(a__     )



       def        __lowercase    (				self      ,    a__     )			->   None:






            '''simple docstring'''





            if item not in self.pos_map:
                 return
            __snake_case      :Dict 							=    self.pos_map[item]
            del self.pos_map[item]
            __snake_case      :Dict 							=    self.arr[self.size - 1]
            __snake_case      :List[str] 							=    index
            self.size -= 1
            # Make sure heap is right in both up and down direction. Ideally only one
            # of them will make any change- so no performance loss in calling both.
            if self.size > index:
                 self._heapify_up(a__     )
                 self._heapify_down(a__     )



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   None:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    len(self.arr     )
            if arr_len == self.size:
                 self.arr.append([item, self.key(a__     )]     )
            else:
                 __snake_case      :Optional[Any] 							=    [item, self.key(a__     )]
            __snake_case      :str 							=    self.size
            self.size += 1
            self._heapify_up(self.size - 1     )



       def        __lowercase    (				self     )			->   tuple | None:






            '''simple docstring'''





            return self.arr[0] if self.size else None



       def        __lowercase    (				self     )			->   tuple | None:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    self.get_top()
            if top_item_tuple:
                 self.delete_item(top_item_tuple[0]     )
            return top_item_tuple






def 	UpperCamelCase	(							):




     '''simple docstring'''








if __name__ == "__main__":
     import doctest

     doctest.testmod()


