





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   lowercase_   (  _A					):
      UpperCamelCase_      :		Optional[Any]										=    "new-model"


if is_tf_available():



 class   lowercase_   (  _A					):
       UpperCamelCase_      :		Optional[Any]										=    NewModelConfig



@require_tf
class   lowercase_   (  unittest.TestCase					):



      @slow
      def 						UpperCamelCase_		(      self				:    Any		)   ->					List[str]:
        _snake_case						=							'''bert-base-cased'''
        _snake_case						=							AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE		)
        self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)
        self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)

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



      @slow
      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					List[str]:
        _snake_case						=							'''bert-base-cased'''
        _snake_case						=							AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE		)
        self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)
        self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)

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



      @slow
      def 						UpperCamelCase_		(      self				:    Tuple		)   ->					Dict:
        for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          _snake_case						=							AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE		)
          self.assertIsNotNone(__SCREAMING_SNAKE_CASE		)
          self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)

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



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

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



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

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



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

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



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

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



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

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



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

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



      def 						UpperCamelCase_		(      self				:    Optional[Any]		)   ->					Optional[Any]:
        _snake_case						=							TFAutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE		)
        self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)
        self.assertEqual(model.num_parameters()      , 14410		)
        self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE		)      , 14410		)



      def 						UpperCamelCase_		(      self				:    Any		)   ->					int:
        _snake_case						=							TFAutoModelWithLMHead.from_pretrained(__SCREAMING_SNAKE_CASE		)
        self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)
        self.assertEqual(model.num_parameters()      , 14410		)
        self.assertEqual(model.num_parameters(only_trainable=__SCREAMING_SNAKE_CASE		)      , 14410		)



      def 						UpperCamelCase_		(      self				:    Tuple		)   ->					Union[str, Any]:
        # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
        _snake_case						=							TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny'''		)
        self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_SNAKE_CASE		)

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

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



      def 						UpperCamelCase_		(      self				:    List[str]		)   ->					List[Any]:
        try:
          AutoConfig.register('''new-model'''      , __SCREAMING_SNAKE_CASE		)

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

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

              with tempfile.TemporaryDirectory() as tmp_dir:
                model.save_pretrained(__SCREAMING_SNAKE_CASE		)
                _snake_case						=							auto_class.from_pretrained(__SCREAMING_SNAKE_CASE		)
                self.assertIsInstance(__SCREAMING_SNAKE_CASE      , __SCREAMING_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 						UpperCamelCase_		(      self				:    List[Any]		)   ->					List[Any]:
        with self.assertRaisesRegex(
            __SCREAMING_SNAKE_CASE      , '''bert-base is not a local folder and is not a valid model identifier'''		):
          _snake_case						=							TFAutoModel.from_pretrained('''bert-base'''		)



      def 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					List[str]:
        with self.assertRaisesRegex(
            __SCREAMING_SNAKE_CASE      , R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''		):
          _snake_case						=							TFAutoModel.from_pretrained(__SCREAMING_SNAKE_CASE      , revision='''aaaaaa'''		)



      def 						UpperCamelCase_		(      self				:    List[str]		)   ->					Optional[Any]:
        with self.assertRaisesRegex(
            __SCREAMING_SNAKE_CASE      , '''hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'''      , ):
          _snake_case						=							TFAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''		)



      def 						UpperCamelCase_		(      self				:    Any		)   ->					Optional[int]:
        with self.assertRaisesRegex(__SCREAMING_SNAKE_CASE      , '''Use `from_pt=True` to load this model'''		):
          _snake_case						=							TFAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''		)







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




import argparse

import torch

from ...utils import logging
from . import AlbertConfig, AlbertForPreTraining, load_tf_weights_in_albert


logging.set_verbosity_info()



def 				snake_case_(_UpperCamelCase			,	_UpperCamelCase			,	_UpperCamelCase					)							->     Optional[Any]:



  """simple docstring"""






  _snake_case						=							AlbertConfig.from_json_file(_UpperCamelCase					)
  print(F"""Building PyTorch model from configuration: {config}"""					)
  _snake_case						=							AlbertForPreTraining(_UpperCamelCase					)

  # Load weights from tf checkpoint
  load_tf_weights_in_albert(_UpperCamelCase			,	_UpperCamelCase			,	_UpperCamelCase					)

  # Save pytorch-model
  print(F"""Save PyTorch model to {pytorch_dump_path}"""					)
  torch.save(model.state_dict()			,	_UpperCamelCase					)


if __name__ == "__main__":
 __A      		=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
 )
 parser.add_argument(
     '''--albert_config_file''',
     default=None,
     type=str,
     required=True,
     help=(
         '''The config json file corresponding to the pre-trained ALBERT model. \n'''
         '''This specifies the model architecture.'''
     ),
 )
 parser.add_argument(
     '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
 )
 __A      		=     parser.parse_args()
 convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.albert_config_file, args.pytorch_dump_path)

