






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()
__a      :       Optional[int]				  =   logging.get_logger(__name__)

__a      :       List[Any]				  =   [
    ("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"),
]

__a      :       Any				  =   [
    "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 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Dict    )						->      List[str]:





    """simple docstring"""

    __A         =	torch.load(__lowercase			,				map_location="""cpu"""    )
    return sd



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Dict			,				__lowercase   :			Dict			,				__lowercase   :			List[Any]=rename_keys_prefix    )						->      Dict:





    """simple docstring"""

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



@torch.no_grad()
def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str			,				__lowercase   :			Tuple    )						->      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:
        __A         =	"""pretraining"""
        if "vcr" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 5_1_2}
        elif "vqa_advanced" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 2_0_4_8}
        elif "vqa" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 2_0_4_8}
        elif "nlvr" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 1_0_2_4}
        else:
            raise NotImplementedError(f"No implementation found for `{checkpoint_path}`."    )
    else:
        if "vcr" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 5_1_2}
            __A         =	"""multichoice"""
        elif "vqa_advanced" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 2_0_4_8}
            __A         =	"""vqa_advanced"""
        elif "vqa" in checkpoint_path:
            __A         =	{"""visual_embedding_dim""": 2_0_4_8, """num_labels""": 3_1_2_9}
            __A         =	"""vqa"""
        elif "nlvr" in checkpoint_path:
            __A         =	{
                """visual_embedding_dim""": 1_0_2_4,
                """num_labels""": 2,
            }
            __A         =	"""nlvr"""

    __A         =	VisualBertConfig(**__lowercase    )

    # Load State Dict
    __A         =	load_state_dict(__lowercase    )

    __A         =	get_new_dict(__lowercase			,				__lowercase    )

    if model_type == "pretraining":
        __A         =	VisualBertForPreTraining(__lowercase    )
    elif model_type == "vqa":
        __A         =	VisualBertForQuestionAnswering(__lowercase    )
    elif model_type == "nlvr":
        __A         =	VisualBertForVisualReasoning(__lowercase    )
    elif model_type == "multichoice":
        __A         =	VisualBertForMultipleChoice(__lowercase    )

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


if __name__ == "__main__":
  __a      :       Dict				  =   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.")
  __a      :       Any				  =   parser.parse_args()
  convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)












import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]    )						->      List[str]:





    """simple docstring"""

    return x + 2






class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)

         __A         =	"""x = y"""
         __A         =	{"""y""": 5}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 5, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""y = add_two(x)"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)

         # Won't work without the tool
         with CaptureStdout() as out:
             __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result is None
         assert "tried to execute add_two" in out.out




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)




     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         __A         =	"""x = 3\ny = 5"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""text = f'This is x: {x}.'"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == "This is x: 3."
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """text""": """This is x: 3."""}				)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	"""if x <= 3:\n    y = 2\nelse:\n    y = 5"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 2}				)

         __A         =	{"""x""": 8}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 8, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertListEqual(UpperCamelCase_ ,	[3, 5]				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	"""y = x"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 3}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]\ntest_list[1]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)

         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)





     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""x = 0\nfor i in range(3):\n    x = i"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{"""range""": range} ,	state=UpperCamelCase_				)
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 2, """i""": 2}				)




