


import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
						import jax

						from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
						from transformers.models.bert.modeling_flax_bert import FlaxBertModel
						from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel
@require_flax
class        _lowercase      (			unittest.TestCase       ):






				'''simple docstring'''



				@slow
				def 		__magic_name__(       self			:str	)		->			int:
										for model_name in ["bert-base-cased", "bert-large-uncased"]:
																with self.subTest(lowerCAmelCase__	):
																						__SCREAMING_SNAKE_CASE						:     Dict					=		AutoConfig.from_pretrained(lowerCAmelCase__	)
																						self.assertIsNotNone(lowerCAmelCase__	)
																						self.assertIsInstance(lowerCAmelCase__			,			lowerCAmelCase__	)

																						__SCREAMING_SNAKE_CASE						:     int					=		FlaxAutoModel.from_pretrained(lowerCAmelCase__	)
																						self.assertIsNotNone(lowerCAmelCase__	)
																						self.assertIsInstance(lowerCAmelCase__			,			lowerCAmelCase__	)


				@slow
				def 		__magic_name__(       self			:Dict	)		->			Dict:
										for model_name in ["roberta-base", "roberta-large"]:
																with self.subTest(lowerCAmelCase__	):
																						__SCREAMING_SNAKE_CASE						:     List[Any]					=		AutoConfig.from_pretrained(lowerCAmelCase__	)
																						self.assertIsNotNone(lowerCAmelCase__	)
																						self.assertIsInstance(lowerCAmelCase__			,			lowerCAmelCase__	)

																						__SCREAMING_SNAKE_CASE						:     Optional[int]					=		FlaxAutoModel.from_pretrained(lowerCAmelCase__	)
																						self.assertIsNotNone(lowerCAmelCase__	)
																						self.assertIsInstance(lowerCAmelCase__			,			lowerCAmelCase__	)


				@slow
				def 		__magic_name__(       self			:Optional[int]	)		->			List[Any]:
										for model_name in ["bert-base-cased", "bert-large-uncased"]:
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		AutoTokenizer.from_pretrained(lowerCAmelCase__	)
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		FlaxBertModel.from_pretrained(lowerCAmelCase__	)
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		tokenizer('''Do you support jax jitted function?'''			,			return_tensors=TensorType.JAX	)

																@jax.jit
																def eval(**lowerCAmelCase__			:List[Any]	):
																						return model(**lowerCAmelCase__	)

																eval(**lowerCAmelCase__	).block_until_ready()


				@slow
				def 		__magic_name__(       self			:List[str]	)		->			Optional[int]:
										for model_name in ["roberta-base", "roberta-large"]:
																__SCREAMING_SNAKE_CASE						:     str					=		AutoTokenizer.from_pretrained(lowerCAmelCase__	)
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		FlaxRobertaModel.from_pretrained(lowerCAmelCase__	)
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		tokenizer('''Do you support jax jitted function?'''			,			return_tensors=TensorType.JAX	)

																@jax.jit
																def eval(**lowerCAmelCase__			:List[Any]	):
																						return model(**lowerCAmelCase__	)

																eval(**lowerCAmelCase__	).block_until_ready()


				def 		__magic_name__(       self			:Dict	)		->			Union[str, Any]:
										with self.assertRaisesRegex(
										    lowerCAmelCase__			,			'''bert-base is not a local folder and is not a valid model identifier'''	):
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		FlaxAutoModel.from_pretrained('''bert-base'''	)


				def 		__magic_name__(       self			:int	)		->			Optional[int]:
										with self.assertRaisesRegex(
										    lowerCAmelCase__			,			r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''	):
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		FlaxAutoModel.from_pretrained(lowerCAmelCase__			,			revision='''aaaaaa'''	)


				def 		__magic_name__(       self			:str	)		->			Tuple:
										with self.assertRaisesRegex(
										    lowerCAmelCase__			,			'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'''			,			):
																__SCREAMING_SNAKE_CASE						:     str					=		FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''	)



				def 		__magic_name__(       self			:Dict	)		->			int:
										with self.assertRaisesRegex(lowerCAmelCase__			,			'''Use `from_pt=True` to load this model'''	):
																__SCREAMING_SNAKE_CASE						:     Dict					=		FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''	)





'''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()
A_							: Union[str, Any]						 =	logging.get_logger(__name__)

A_							: Union[str, 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_							: int						 =	[
    """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_					(     lowerCAmelCase_		)->		Optional[int]:





       '''simple docstring'''


       _UpperCAmelCase  :						Dict			     = torch.load(lowerCAmelCase_       ,   map_location="""cpu"""		)
       return sd




def     snake_case_					(     lowerCAmelCase_       ,   lowerCAmelCase_       ,   lowerCAmelCase_=rename_keys_prefix		)->		Optional[Any]:





       '''simple docstring'''


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




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

       _UpperCAmelCase  :						int			     = VisualBertConfig(**lowerCAmelCase_		)

       # Load State Dict
       _UpperCAmelCase  :						Any			     = load_state_dict(lowerCAmelCase_		)

       _UpperCAmelCase  :						Optional[int]			     = get_new_dict(lowerCAmelCase_       ,   lowerCAmelCase_		)

       if model_type == "pretraining":
              _UpperCAmelCase  :						List[str]			     = VisualBertForPreTraining(lowerCAmelCase_		)
       elif model_type == "vqa":
              _UpperCAmelCase  :						Optional[int]			     = VisualBertForQuestionAnswering(lowerCAmelCase_		)
       elif model_type == "nlvr":
              _UpperCAmelCase  :						str			     = VisualBertForVisualReasoning(lowerCAmelCase_		)
       elif model_type == "multichoice":
              _UpperCAmelCase  :						Dict			     = VisualBertForMultipleChoice(lowerCAmelCase_		)

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


if __name__ == "__main__":
   A_							: Optional[Any]						 =	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_							: int						 =	parser.parse_args()
   convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)



