


"""simple docstring"""


from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
	from transformers import (
	    TFAutoModel,
	    TFAutoModelForCausalLM,
	    TFAutoModelForMaskedLM,
	    TFAutoModelForPreTraining,
	    TFAutoModelForQuestionAnswering,
	    TFAutoModelForSeqaSeqLM,
	    TFAutoModelForSequenceClassification,
	    TFAutoModelForTableQuestionAnswering,
	    TFAutoModelForTokenClassification,
	    TFAutoModelWithLMHead,
	    TFBertForMaskedLM,
	    TFBertForPreTraining,
	    TFBertForQuestionAnswering,
	    TFBertForSequenceClassification,
	    TFBertModel,
	    TFFunnelBaseModel,
	    TFFunnelModel,
	    TFGPTaLMHeadModel,
	    TFRobertaForMaskedLM,
	    TFTaForConditionalGeneration,
	    TFTapasForQuestionAnswering,
	)
	from transformers.models.auto.modeling_tf_auto import (
	    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
	    TF_MODEL_FOR_MASKED_LM_MAPPING,
	    TF_MODEL_FOR_PRETRAINING_MAPPING,
	    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
	    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
	    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
	    TF_MODEL_MAPPING,
	)
	from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
	from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
	from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
	from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST
class   _UpperCamelCase  (   lowerCAmelCase__		):
							'''simple docstring'''







							__UpperCAmelCase						:    Tuple				      ="""new-model"""


if is_tf_available():
	class   _UpperCamelCase  (   lowerCAmelCase__		):
								'''simple docstring'''







								__UpperCAmelCase						:    Optional[int]				      =NewModelConfig
@require_tf
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







							@slow
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			"bert-base-cased"
										__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
										self.assertIsNotNone(__a	)
										self.assertIsInstance(__a    ,   __a	)

										__lowerCAmelCase				    =			TFAutoModel.from_pretrained(__a	)
										self.assertIsNotNone(__a	)
										self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			"bert-base-cased"
										__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
										self.assertIsNotNone(__a	)
										self.assertIsInstance(__a    ,   __a	)

										__lowerCAmelCase				    =			TFAutoModelForPreTraining.from_pretrained(__a	)
										self.assertIsNotNone(__a	)
										self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForCausalLM.from_pretrained(__a	)
													__lowerCAmelCase    ,	__lowerCAmelCase				    =			TFAutoModelForCausalLM.from_pretrained(__a    ,   output_loading_info=__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelWithLMHead.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForMaskedLM.from_pretrained(__a	)
													__lowerCAmelCase    ,	__lowerCAmelCase				    =			TFAutoModelForMaskedLM.from_pretrained(__a    ,   output_loading_info=__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForSeqaSeqLM.from_pretrained(__a	)
													__lowerCAmelCase    ,	__lowerCAmelCase				    =			TFAutoModelForSeqaSeqLM.from_pretrained(__a    ,   output_loading_info=__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForSequenceClassification.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							def 			snake_case				(     self	):
										# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										for model_name in ["bert-base-uncased"]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForQuestionAnswering.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							@slow
							@require_tensorflow_probability
							def 			snake_case				(     self	):
										for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
													__lowerCAmelCase				    =			AutoConfig.from_pretrained(__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

													__lowerCAmelCase				    =			TFAutoModelForTableQuestionAnswering.from_pretrained(__a	)
													__lowerCAmelCase    ,	__lowerCAmelCase				    =			TFAutoModelForTableQuestionAnswering.from_pretrained(
													    __a    ,   output_loading_info=__a	)
													self.assertIsNotNone(__a	)
													self.assertIsInstance(__a    ,   __a	)

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			TFAutoModelWithLMHead.from_pretrained(__a	)
										self.assertIsInstance(__a    ,   __a	)
										self.assertEqual(model.num_parameters()    ,   1_44_10	)
										self.assertEqual(model.num_parameters(only_trainable=__a	)    ,   1_44_10	)

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			TFAutoModelWithLMHead.from_pretrained(__a	)
										self.assertIsInstance(__a    ,   __a	)
										self.assertEqual(model.num_parameters()    ,   1_44_10	)
										self.assertEqual(model.num_parameters(only_trainable=__a	)    ,   1_44_10	)

							def 			snake_case				(     self	):
										# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
										__lowerCAmelCase				    =			TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"	)
										self.assertIsInstance(__a    ,   __a	)

										__lowerCAmelCase				    =			copy.deepcopy(model.config	)
										__lowerCAmelCase				    =			["FunnelBaseModel"]
										__lowerCAmelCase				    =			TFAutoModel.from_config(__a	)
										self.assertIsInstance(__a    ,   __a	)

										with tempfile.TemporaryDirectory() as tmp_dir:
													model.save_pretrained(__a	)
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained(__a	)
													self.assertIsInstance(__a    ,   __a	)

							def 			snake_case				(     self	):
										try:
													AutoConfig.register("new-model"    ,   __a	)

													__lowerCAmelCase				    =			[
													    TFAutoModel,
													    TFAutoModelForCausalLM,
													    TFAutoModelForMaskedLM,
													    TFAutoModelForPreTraining,
													    TFAutoModelForQuestionAnswering,
													    TFAutoModelForSequenceClassification,
													    TFAutoModelForTokenClassification,
													]

													for auto_class in auto_classes:
																with self.subTest(auto_class.__name__	):
																			# Wrong config class will raise an error
																			with self.assertRaises(__a	):
																						auto_class.register(__a    ,   __a	)
																			auto_class.register(__a    ,   __a	)
																			# Trying to register something existing in the Transformers library will raise an error
																			with self.assertRaises(__a	):
																						auto_class.register(__a    ,   __a	)

																			# Now that the config is registered, it can be used as any other config with the auto-API
																			__lowerCAmelCase				    =			BertModelTester(self	).get_config()
																			__lowerCAmelCase				    =			NewModelConfig(**tiny_config.to_dict()	)
																			__lowerCAmelCase				    =			auto_class.from_config(__a	)
																			self.assertIsInstance(__a    ,   __a	)

																			with tempfile.TemporaryDirectory() as tmp_dir:
																						model.save_pretrained(__a	)
																						__lowerCAmelCase				    =			auto_class.from_pretrained(__a	)
																						self.assertIsInstance(__a    ,   __a	)

										finally:
													if "new-model" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["new-model"]
													for mapping in (
													    TF_MODEL_MAPPING,
													    TF_MODEL_FOR_PRETRAINING_MAPPING,
													    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
													    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
													    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
													    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
													    TF_MODEL_FOR_MASKED_LM_MAPPING,
													):
																if NewModelConfig in mapping._extra_content:
																			del mapping._extra_content[NewModelConfig]

							def 			snake_case				(     self	):
										with self.assertRaisesRegex(
										    __a    ,   "bert-base is not a local folder and is not a valid model identifier"	):
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained("bert-base"	)

							def 			snake_case				(     self	):
										with self.assertRaisesRegex(
										    __a    ,   R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"	):
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained(__a    ,   revision="aaaaaa"	)

							def 			snake_case				(     self	):
										with self.assertRaisesRegex(
										    __a    ,   "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"    ,   ):
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained("hf-internal-testing/config-no-model"	)

							def 			snake_case				(     self	):
										with self.assertRaisesRegex(__a    ,   "Use `from_pt=True` to load this model"	):
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"	)



							def 			snake_case				(     self	):
										# Make sure we have cached the model.
										__lowerCAmelCase				    =			TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"	)
										with RequestCounter() as counter:
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"	)
													self.assertEqual(counter.get_request_count    ,   0	)
													self.assertEqual(counter.head_request_count    ,   1	)
													self.assertEqual(counter.other_request_count    ,   0	)

										# With a sharded checkpoint
										__lowerCAmelCase				    =			TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"	)
										with RequestCounter() as counter:
													__lowerCAmelCase				    =			TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"	)
													self.assertEqual(counter.get_request_count    ,   0	)
													self.assertEqual(counter.head_request_count    ,   1	)
													self.assertEqual(counter.other_request_count    ,   0	)








"""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     :   Tuple			    =			logging.get_logger(__name__)

A     :   Tuple			    =			[
    ("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     :   Optional[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 _lowerCamelCase	(		_UpperCamelCase      ):

			'''simple docstring'''




			__lowerCAmelCase				    =			torch.load(_UpperCamelCase    , map_location="cpu"      )
			return sd


def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase    , _UpperCamelCase=rename_keys_prefix      ):

			'''simple docstring'''




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


@torch.no_grad()
def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase      ):

			'''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:
						__lowerCAmelCase				    =			"pretraining"
						if "vcr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 512}
						elif "vqa_advanced" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
						elif "vqa" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
						elif "nlvr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 1024}
						else:
									raise NotImplementedError(f"No implementation found for `{checkpoint_path}`."      )
			else:
						if "vcr" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 512}
									__lowerCAmelCase				    =			"multichoice"
						elif "vqa_advanced" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048}
									__lowerCAmelCase				    =			"vqa_advanced"
						elif "vqa" in checkpoint_path:
									__lowerCAmelCase				    =			{"visual_embedding_dim": 2048, "num_labels": 3129}
									__lowerCAmelCase				    =			"vqa"
						elif "nlvr" in checkpoint_path:
									__lowerCAmelCase				    =			{
									    "visual_embedding_dim": 1024,
									    "num_labels": 2,
									}
									__lowerCAmelCase				    =			"nlvr"

			__lowerCAmelCase				    =			VisualBertConfig(**_UpperCamelCase      )

			# Load State Dict
			__lowerCAmelCase				    =			load_state_dict(_UpperCamelCase      )

			__lowerCAmelCase				    =			get_new_dict(_UpperCamelCase    , _UpperCamelCase      )

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

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


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




