"""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 lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






					'''simple docstring'''


					a   	=		"new-model"


if is_tf_available():

				class lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






									'''simple docstring'''


									a   	=		NewModelConfig

@require_tf
class lowerCamelCase		(unittest.TestCase      ):






					'''simple docstring'''


					@slow
					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				"bert-base-cased"
									SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Any		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__		  =				"bert-base-cased"
									SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)

									SCREAMING_SNAKE_CASE__		  =				TFAutoModelForPreTraining.from_pretrained(_snake_case		)
									self.assertIsNotNone(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[Any]:
									for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForCausalLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForCausalLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       str:
									for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Optional[int]:
									for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForMaskedLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForMaskedLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				str		)							->       Any:
									for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSeqaSeqLM.from_pretrained(_snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       Tuple:
									# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									for model_name in ["bert-base-uncased"]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForSequenceClassification.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       Optional[Any]:
									# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									for model_name in ["bert-base-uncased"]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForQuestionAnswering.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					@slow
					@require_tensorflow_probability
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Dict:
									for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
													SCREAMING_SNAKE_CASE__		  =				AutoConfig.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				TFAutoModelForTableQuestionAnswering.from_pretrained(_snake_case		)
													SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				TFAutoModelForTableQuestionAnswering.from_pretrained(
													    _snake_case   , output_loading_info=_snake_case		)
													self.assertIsNotNone(_snake_case		)
													self.assertIsInstance(_snake_case   , _snake_case		)



					def 			lowerCAmelCase_					(						self			:				int		)							->       Optional[int]:
									SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)
									self.assertEqual(model.num_parameters()   , 14410		)
									self.assertEqual(model.num_parameters(only_trainable=_snake_case		)   , 14410		)



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__		  =				TFAutoModelWithLMHead.from_pretrained(_snake_case		)
									self.assertIsInstance(_snake_case   , _snake_case		)
									self.assertEqual(model.num_parameters()   , 14410		)
									self.assertEqual(model.num_parameters(only_trainable=_snake_case		)   , 14410		)



					def 			lowerCAmelCase_					(						self			:				Any		)							->       List[Any]:
									# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"		)
									self.assertIsInstance(_snake_case   , _snake_case		)

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

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



					def 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[Any]:
									try:
													AutoConfig.register("new-model"   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				[
													    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(_snake_case		):
																									auto_class.register(_snake_case   , _snake_case		)
																					auto_class.register(_snake_case   , _snake_case		)
																					# Trying to register something existing in the Transformers library will raise an error
																					with self.assertRaises(_snake_case		):
																									auto_class.register(_snake_case   , _snake_case		)

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

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

									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 			lowerCAmelCase_					(						self			:				Optional[Any]		)							->       Optional[int]:
									with self.assertRaisesRegex(
									    _snake_case   , "bert-base is not a local folder and is not a valid model identifier"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("bert-base"		)



					def 			lowerCAmelCase_					(						self			:				int		)							->       List[str]:
									with self.assertRaisesRegex(
									    _snake_case   , r"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained(_snake_case   , revision="aaaaaa"		)



					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       int:
									with self.assertRaisesRegex(
									    _snake_case   , "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"   , ):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/config-no-model"		)



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       Tuple:
									with self.assertRaisesRegex(_snake_case   , "Use `from_pt=True` to load this model"		):
													SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"		)



					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Optional[Any]:
									# Make sure we have cached the model.
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"		)
									with RequestCounter() as counter:
													SCREAMING_SNAKE_CASE__		  =				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
									SCREAMING_SNAKE_CASE__		  =				TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"		)
									with RequestCounter() as counter:
													SCREAMING_SNAKE_CASE__		  =				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
import os

import torch

from transformers.utils import WEIGHTS_NAME


_A   					=    ['small', 'medium', 'large']

_A   					=    'lm_head.decoder.weight'
_A   					=    'lm_head.weight'




def 						SCREAMING_SNAKE_CASE					(  __UpperCAmelCase  ,					__UpperCAmelCase       )							->   List[str]:
				SCREAMING_SNAKE_CASE__		  =				torch.load(__UpperCAmelCase       )
				SCREAMING_SNAKE_CASE__		  =				d.pop(__UpperCAmelCase       )
				os.makedirs(__UpperCAmelCase  ,					exist_ok=__UpperCAmelCase       )
				torch.save(__UpperCAmelCase  ,					os.path.join(__UpperCAmelCase  ,					__UpperCAmelCase       )       )


if __name__ == "__main__":
						_A   					=    argparse.ArgumentParser()
						parser.add_argument('--dialogpt_path', default='.', type=str)
						_A   					=    parser.parse_args()
						for MODEL in DIALOGPT_MODELS:
												_A   					=    os.path.join(args.dialogpt_path, F'{MODEL}_ft.pkl')
												_A   					=    F'./DialoGPT-{MODEL}'
												convert_dialogpt_checkpoint(
												    checkpoint_path,
												    pytorch_dump_folder_path,
												)




