




'''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_				(	a__					):
 UpperCAmelCase_		    =						"new-model"


if is_tf_available():






  class    UpperCAmelCase_				(	a__					):
   UpperCAmelCase_		    =						NewModelConfig






@require_tf
class    UpperCAmelCase_				(	unittest.TestCase					):





 @slow
 def        snake_case__		(    self):
     snake_case_ :    List[Any]     			=						"bert-base-cased"
     snake_case_ :    Union[str, Any]     			=						AutoConfig.from_pretrained(lowerCamelCase_)
     self.assertIsNotNone(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

     snake_case_ :    Optional[Any]     			=						TFAutoModel.from_pretrained(lowerCamelCase_)
     self.assertIsNotNone(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     snake_case_ :    Tuple     			=						"bert-base-cased"
     snake_case_ :    Dict     			=						AutoConfig.from_pretrained(lowerCamelCase_)
     self.assertIsNotNone(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

     snake_case_ :    Tuple     			=						TFAutoModelForPreTraining.from_pretrained(lowerCamelCase_)
     self.assertIsNotNone(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    Union[str, Any]     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    Union[str, Any]     			=						TFAutoModelForCausalLM.from_pretrained(lowerCamelCase_)
         snake_case_			, snake_case_ :    Tuple     			=						TFAutoModelForCausalLM.from_pretrained(lowerCamelCase_ ,       output_loading_info=lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    Optional[Any]     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    List[str]     			=						TFAutoModelWithLMHead.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    int     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    List[str]     			=						TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase_)
         snake_case_			, snake_case_ :    List[Any]     			=						TFAutoModelForMaskedLM.from_pretrained(lowerCamelCase_ ,       output_loading_info=lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :    List[Any]     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    Any     			=						TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase_)
         snake_case_			, snake_case_ :    Optional[Any]     			=						TFAutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase_ ,       output_loading_info=lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





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

         snake_case_ :    List[Any]     			=						TFAutoModelForSequenceClassification.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 def        snake_case__		(    self):
     # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     for model_name in ["bert-base-uncased"]:
         snake_case_ :    Dict     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    Dict     			=						TFAutoModelForQuestionAnswering.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 @slow
 @require_tensorflow_probability
 def        snake_case__		(    self):
     for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
         snake_case_ :    str     			=						AutoConfig.from_pretrained(lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

         snake_case_ :    Dict     			=						TFAutoModelForTableQuestionAnswering.from_pretrained(lowerCamelCase_)
         snake_case_			, snake_case_ :    List[str]     			=						TFAutoModelForTableQuestionAnswering.from_pretrained(
             lowerCamelCase_ ,       output_loading_info=lowerCamelCase_)
         self.assertIsNotNone(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 def        snake_case__		(    self):
     snake_case_ :    List[str]     			=						TFAutoModelWithLMHead.from_pretrained(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)
     self.assertEqual(model.num_parameters() ,       1_44_10)
     self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase_) ,       1_44_10)





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						TFAutoModelWithLMHead.from_pretrained(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)
     self.assertEqual(model.num_parameters() ,       1_44_10)
     self.assertEqual(model.num_parameters(only_trainable=lowerCamelCase_) ,       1_44_10)





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

     snake_case_ :    List[Any]     			=						copy.deepcopy(model.config)
     snake_case_ :    str     			=						["FunnelBaseModel"]
     snake_case_ :    List[Any]     			=						TFAutoModel.from_config(lowerCamelCase_)
     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

     with tempfile.TemporaryDirectory() as tmp_dir:
         model.save_pretrained(lowerCamelCase_)
         snake_case_ :    Tuple     			=						TFAutoModel.from_pretrained(lowerCamelCase_)
         self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)





 def        snake_case__		(    self):
     try:
         AutoConfig.register("new-model" ,       lowerCamelCase_)

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

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

                 with tempfile.TemporaryDirectory() as tmp_dir:
                     model.save_pretrained(lowerCamelCase_)
                     snake_case_ :    Optional[int]     			=						auto_class.from_pretrained(lowerCamelCase_)
                     self.assertIsInstance(lowerCamelCase_ ,       lowerCamelCase_)

     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(
         lowerCamelCase_ ,       "bert-base is not a local folder and is not a valid model identifier"):
         snake_case_ :    Optional[int]     			=						TFAutoModel.from_pretrained("bert-base")





 def        snake_case__		(    self):
     with self.assertRaisesRegex(
         lowerCamelCase_ ,       R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"):
         snake_case_ :    Optional[int]     			=						TFAutoModel.from_pretrained(lowerCamelCase_ ,       revision="aaaaaa")





 def        snake_case__		(    self):
     with self.assertRaisesRegex(
         lowerCamelCase_ ,       "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin" ,       ):
         snake_case_ :    Optional[int]     			=						TFAutoModel.from_pretrained("hf-internal-testing/config-no-model")





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





 def        snake_case__		(    self):
     # Make sure we have cached the model.
     snake_case_ :    str     			=						TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert")
     with RequestCounter() as counter:
         snake_case_ :    int     			=						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
     snake_case_ :    str     			=						TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
     with RequestCounter() as counter:
         snake_case_ :    str     			=						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 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    List[Any]     			=						torch.load(__SCREAMING_SNAKE_CASE       )
    snake_case_ :    List[Any]     			=						d.pop(__SCREAMING_SNAKE_CASE       )
    os.makedirs(__SCREAMING_SNAKE_CASE,		exist_ok=__SCREAMING_SNAKE_CASE       )
    torch.save(__SCREAMING_SNAKE_CASE,		os.path.join(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )       )


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,
    )




