






'''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 typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
 from PIL import Image

 from ..image_utils import load_image

if is_tf_available():
 import tensorflow as tf

 from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
 from ..tf_utils import stable_softmax

if is_torch_available():
 from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

lowercase				=logging.get_logger(__name__)






@add_end_docstrings(lowerCAmelCase							)
class    __magic_name__				(						lowerCAmelCase							):






       def __init__(     self    ,				*snake_case    ,				**snake_case)		->    Dict:




          '''simple docstring'''
          super().__init__(*snake_case    ,				**snake_case)
          requires_backends(self    ,				'vision')
          self.check_model_type(
              TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
              if self.framework == 'tf'
              else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING)






       def      lowerCAmelCase						(     self    ,				snake_case=None)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       ={}
          if top_k is not None:
             _UpperCAmelCase  :  Optional[int]				       =top_k
          return {}, {}, postprocess_params






       def __call__(     self    ,				snake_case    ,				**snake_case)		->    List[str]:




          '''simple docstring'''
          return super().__call__(snake_case    ,				**snake_case)






       def      lowerCAmelCase						(     self    ,				snake_case)		->    List[str]:




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =load_image(snake_case)
          _UpperCAmelCase  :  Tuple				       =self.image_processor(images=snake_case    ,				return_tensors=self.framework)
          return model_inputs






       def      lowerCAmelCase						(     self    ,				snake_case)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =self.model(**snake_case)
          return model_outputs






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case=5)		->    Any:




          '''simple docstring'''
          if top_k > self.model.config.num_labels:
             _UpperCAmelCase  :  Optional[Any]				       =self.model.config.num_labels

          if self.framework == "pt":
             _UpperCAmelCase  :  List[Any]				       =model_outputs.logits.softmax(-1)[0]
             _UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[Any]				       =probs.topk(snake_case)
          elif self.framework == "tf":
             _UpperCAmelCase  :  int				       =stable_softmax(model_outputs.logits    ,				axis=-1)[0]
             _UpperCAmelCase  :  List[Any]				       =tf.math.top_k(snake_case    ,				k=snake_case)
             _UpperCAmelCase       ,			_UpperCAmelCase  :  List[str]				       =topk.values.numpy(), topk.indices.numpy()
          else:
             raise ValueError(f"Unsupported framework: {self.framework}")

          _UpperCAmelCase  :  int				       =scores.tolist()
          _UpperCAmelCase  :  Tuple				       =ids.tolist()
          return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(snake_case    ,				snake_case)]

