



"""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()
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      a__			(   lowerCAmelCase__   ):
      UpperCAmelCase_    				=							torch.load(snake_case__			,    map_location="cpu"   )
      return sd
def      a__			(   lowerCAmelCase__			,    lowerCAmelCase__			,    lowerCAmelCase__=rename_keys_prefix   ):
      UpperCAmelCase_    				=							OrderedDict()
      UpperCAmelCase_    				=							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_    				=							key
            for name_pair in rename_keys_prefix:
                  UpperCAmelCase_    				=							new_key.replace(name_pair[0]			,    name_pair[1]   )
            UpperCAmelCase_    				=							d[key]
            if key == "bert.cls.predictions.decoder.weight":
                  # Old bert code didn't have `decoder.bias`, but was added separately
                  UpperCAmelCase_    				=							new_d["cls.predictions.bias"]
      return new_d
@torch.no_grad()
def      a__			(   lowerCAmelCase__			,    lowerCAmelCase__   ):
      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_    				=							"pretraining"
            if "vcr" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 512}
            elif "vqa_advanced" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 2048}
            elif "vqa" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 2048}
            elif "nlvr" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 1024}
            else:
                  raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""   )
      else:
            if "vcr" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 512}
                  UpperCAmelCase_    				=							"multichoice"
            elif "vqa_advanced" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 2048}
                  UpperCAmelCase_    				=							"vqa_advanced"
            elif "vqa" in checkpoint_path:
                  UpperCAmelCase_    				=							{"visual_embedding_dim": 2048, "num_labels": 3129}
                  UpperCAmelCase_    				=							"vqa"
            elif "nlvr" in checkpoint_path:
                  UpperCAmelCase_    				=							{
                      "visual_embedding_dim": 1024,
                      "num_labels": 2,
                  }
                  UpperCAmelCase_    				=							"nlvr"

      UpperCAmelCase_    				=							VisualBertConfig(**snake_case__   )

      # Load State Dict
      UpperCAmelCase_    				=							load_state_dict(snake_case__   )

      UpperCAmelCase_    				=							get_new_dict(snake_case__			,    snake_case__   )

      if model_type == "pretraining":
            UpperCAmelCase_    				=							VisualBertForPreTraining(snake_case__   )
      elif model_type == "vqa":
            UpperCAmelCase_    				=							VisualBertForQuestionAnswering(snake_case__   )
      elif model_type == "nlvr":
            UpperCAmelCase_    				=							VisualBertForVisualReasoning(snake_case__   )
      elif model_type == "multichoice":
            UpperCAmelCase_    				=							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)









"""simple docstring"""





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
  import PIL


lowerCamelCase      =		logging.get_logger(__name__)
def      a__			(   lowerCAmelCase__   ):
      if isinstance(lowerCAmelCase__			,    (list, tuple)   ) and isinstance(videos[0]			,    (list, tuple)   ) and is_valid_image(videos[0][0]   ):
            return videos

      elif isinstance(lowerCAmelCase__			,    (list, tuple)   ) and is_valid_image(videos[0]   ):
            return [videos]

      elif is_valid_image(lowerCAmelCase__   ):
            return [[videos]]

      raise ValueError(f"""Could not make batched video from {videos}"""   )


class    lowercase__     (					SCREAMING_SNAKE_CASE   ):







       '''simple docstring'''


       UpperCamelCase		=  ['''pixel_values''']



       def __init__(							self    :						Tuple	,       _UpperCAmelCase    :						bool = True	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						PILImageResampling = PILImageResampling.BILINEAR	,       _UpperCAmelCase    :						bool = True	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						bool = True	,       _UpperCAmelCase    :						Union[int, float] = 1 / 255	,       _UpperCAmelCase    :						bool = True	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       **_UpperCAmelCase    :						Any	,       )				->		None:


             '''simple docstring'''




             super().__init__(**_UpperCAmelCase					)
             UpperCAmelCase_    				=							size if size is not None else {"shortest_edge": 224}
             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase	,       default_to_square=_UpperCAmelCase					)
             UpperCAmelCase_    				=							crop_size if crop_size is not None else {"height": 224, "width": 224}
             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase	,       param_name="crop_size"					)

             UpperCAmelCase_    				=							do_resize
             UpperCAmelCase_    				=							size
             UpperCAmelCase_    				=							do_center_crop
             UpperCAmelCase_    				=							crop_size
             UpperCAmelCase_    				=							resample
             UpperCAmelCase_    				=							do_rescale
             UpperCAmelCase_    				=							rescale_factor
             UpperCAmelCase_    				=							do_normalize
             UpperCAmelCase_    				=							image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
             UpperCAmelCase_    				=							image_std if image_std is not None else IMAGENET_STANDARD_STD



       def   lowercase__       (							self    :						Tuple	,       _UpperCAmelCase    :						np.ndarray	,       _UpperCAmelCase    :						Dict[str, int]	,       _UpperCAmelCase    :						PILImageResampling = PILImageResampling.BILINEAR	,       _UpperCAmelCase    :						Optional[Union[str, ChannelDimension]] = None	,       **_UpperCAmelCase    :						Tuple	,       )				->		np.ndarray:


             '''simple docstring'''




             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase	,       default_to_square=_UpperCAmelCase					)
             if "shortest_edge" in size:
                   UpperCAmelCase_    				=							get_resize_output_image_size(_UpperCAmelCase	,       size["shortest_edge"]	,       default_to_square=_UpperCAmelCase					)
             elif "height" in size and "width" in size:
                   UpperCAmelCase_    				=							(size["height"], size["width"])
             else:
                   raise ValueError(F"""Size must have 'height' and 'width' or 'shortest_edge' as keys. Got {size.keys()}"""					)
             return resize(_UpperCAmelCase	,       size=_UpperCAmelCase	,       resample=_UpperCAmelCase	,       data_format=_UpperCAmelCase	,       **_UpperCAmelCase					)



       def   lowercase__       (							self    :						Union[str, Any]	,       _UpperCAmelCase    :						np.ndarray	,       _UpperCAmelCase    :						Dict[str, int]	,       _UpperCAmelCase    :						Optional[Union[str, ChannelDimension]] = None	,       **_UpperCAmelCase    :						Union[str, Any]	,       )				->		np.ndarray:


             '''simple docstring'''




             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase					)
             if "height" not in size or "width" not in size:
                   raise ValueError(F"""Size must have 'height' and 'width' as keys. Got {size.keys()}"""					)
             return center_crop(_UpperCAmelCase	,       size=(size["height"], size["width"])	,       data_format=_UpperCAmelCase	,       **_UpperCAmelCase					)



       def   lowercase__       (							self    :						List[Any]	,       _UpperCAmelCase    :						np.ndarray	,       _UpperCAmelCase    :						Union[int, float]	,       _UpperCAmelCase    :						Optional[Union[str, ChannelDimension]] = None	,       **_UpperCAmelCase    :						str	,       )				->		List[str]:


             '''simple docstring'''




             return rescale(_UpperCAmelCase	,       scale=_UpperCAmelCase	,       data_format=_UpperCAmelCase	,       **_UpperCAmelCase					)



       def   lowercase__       (							self    :						Optional[Any]	,       _UpperCAmelCase    :						np.ndarray	,       _UpperCAmelCase    :						Union[float, List[float]]	,       _UpperCAmelCase    :						Union[float, List[float]]	,       _UpperCAmelCase    :						Optional[Union[str, ChannelDimension]] = None	,       **_UpperCAmelCase    :						Tuple	,       )				->		np.ndarray:


             '''simple docstring'''




             return normalize(_UpperCAmelCase	,       mean=_UpperCAmelCase	,       std=_UpperCAmelCase	,       data_format=_UpperCAmelCase	,       **_UpperCAmelCase					)



       def   lowercase__       (							self    :						Any	,       _UpperCAmelCase    :						ImageInput	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						PILImageResampling = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						float = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       _UpperCAmelCase    :						Optional[ChannelDimension] = ChannelDimension.FIRST	,       )				->		np.ndarray:


             '''simple docstring'''




             if do_resize and size is None or resample is None:
                   raise ValueError("Size and resample must be specified if do_resize is True."					)

             if do_center_crop and crop_size is None:
                   raise ValueError("Crop size must be specified if do_center_crop is True."					)

             if do_rescale and rescale_factor is None:
                   raise ValueError("Rescale factor must be specified if do_rescale is True."					)

             if do_normalize and (image_mean is None or image_std is None):
                   raise ValueError("Image mean and std must be specified if do_normalize is True."					)

             # All transformations expect numpy arrays.
             UpperCAmelCase_    				=							to_numpy_array(_UpperCAmelCase					)

             if do_resize:
                   UpperCAmelCase_    				=							self.resize(image=_UpperCAmelCase	,       size=_UpperCAmelCase	,       resample=_UpperCAmelCase					)

             if do_center_crop:
                   UpperCAmelCase_    				=							self.center_crop(_UpperCAmelCase	,       size=_UpperCAmelCase					)

             if do_rescale:
                   UpperCAmelCase_    				=							self.rescale(image=_UpperCAmelCase	,       scale=_UpperCAmelCase					)

             if do_normalize:
                   UpperCAmelCase_    				=							self.normalize(image=_UpperCAmelCase	,       mean=_UpperCAmelCase	,       std=_UpperCAmelCase					)

             UpperCAmelCase_    				=							to_channel_dimension_format(_UpperCAmelCase	,       _UpperCAmelCase					)
             return image







       def   lowercase__       (							self    :						Optional[Any]	,       _UpperCAmelCase    :						ImageInput	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						PILImageResampling = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Dict[str, int] = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						float = None	,       _UpperCAmelCase    :						bool = None	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       _UpperCAmelCase    :						Optional[Union[float, List[float]]] = None	,       _UpperCAmelCase    :						Optional[Union[str, TensorType]] = None	,       _UpperCAmelCase    :						ChannelDimension = ChannelDimension.FIRST	,       **_UpperCAmelCase    :						int	,       )				->		PIL.Image.Image:


             '''simple docstring'''




             UpperCAmelCase_    				=							do_resize if do_resize is not None else self.do_resize
             UpperCAmelCase_    				=							resample if resample is not None else self.resample
             UpperCAmelCase_    				=							do_center_crop if do_center_crop is not None else self.do_center_crop
             UpperCAmelCase_    				=							do_rescale if do_rescale is not None else self.do_rescale
             UpperCAmelCase_    				=							rescale_factor if rescale_factor is not None else self.rescale_factor
             UpperCAmelCase_    				=							do_normalize if do_normalize is not None else self.do_normalize
             UpperCAmelCase_    				=							image_mean if image_mean is not None else self.image_mean
             UpperCAmelCase_    				=							image_std if image_std is not None else self.image_std

             UpperCAmelCase_    				=							size if size is not None else self.size
             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase	,       default_to_square=_UpperCAmelCase					)
             UpperCAmelCase_    				=							crop_size if crop_size is not None else self.crop_size
             UpperCAmelCase_    				=							get_size_dict(_UpperCAmelCase	,       param_name="crop_size"					)

             if not valid_images(_UpperCAmelCase					):
                   raise ValueError(
                       "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                       "torch.Tensor, tf.Tensor or jax.ndarray."					)

             UpperCAmelCase_    				=							make_batched(_UpperCAmelCase					)

             UpperCAmelCase_    				=							[
                 [
                     self._preprocess_image(
                         image=_UpperCAmelCase	,       do_resize=_UpperCAmelCase	,       size=_UpperCAmelCase	,       resample=_UpperCAmelCase	,       do_center_crop=_UpperCAmelCase	,       crop_size=_UpperCAmelCase	,       do_rescale=_UpperCAmelCase	,       rescale_factor=_UpperCAmelCase	,       do_normalize=_UpperCAmelCase	,       image_mean=_UpperCAmelCase	,       image_std=_UpperCAmelCase	,       data_format=_UpperCAmelCase	,       )
                     for img in video
                 ]
                 for video in videos
             ]

             UpperCAmelCase_    				=							{"pixel_values": videos}
             return BatchFeature(data=_UpperCAmelCase	,       tensor_type=_UpperCAmelCase					)




