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__    (		snake_case					):



  """simple docstring"""






  lowerCAmelCase__		:	str    		= 'new-model'


if is_tf_available():


       class 			UpperCAmelCase__    (		snake_case					):



         """simple docstring"""






         lowerCAmelCase__		:	Union[str, Any]    		= NewModelConfig


@require_tf
class 			UpperCAmelCase__    (		unittest.TestCase					):



  """simple docstring"""






  @slow
  def 	_UpperCAmelCase				(				self:			int   ) -> int:




    '''simple docstring'''




    __UpperCAmelCase						    =       "bert-base-cased"
    __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
    self.assertIsNotNone(__lowerCAmelCase   )
    self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Optional[int]   ) -> Optional[int]:




    '''simple docstring'''




    __UpperCAmelCase						    =       "bert-base-cased"
    __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
    self.assertIsNotNone(__lowerCAmelCase   )
    self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Any   ) -> Optional[int]:




    '''simple docstring'''




    for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Tuple   ) -> Tuple:




    '''simple docstring'''




    for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Tuple   ) -> Optional[Any]:




    '''simple docstring'''




    for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			List[Any]   ) -> Any:




    '''simple docstring'''




    for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Any   ) -> int:




    '''simple docstring'''




    for model_name in ["bert-base-uncased"]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  def 	_UpperCAmelCase				(				self:			Union[str, Any]   ) -> List[Any]:




    '''simple docstring'''




    for model_name in ["bert-base-uncased"]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  @slow
  @require_tensorflow_probability
  def 	_UpperCAmelCase				(				self:			int   ) -> str:




    '''simple docstring'''




    for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
      __UpperCAmelCase						    =       AutoConfig.from_pretrained(__lowerCAmelCase   )
      self.assertIsNotNone(__lowerCAmelCase   )
      self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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






  def 	_UpperCAmelCase				(				self:			Any   ) -> Dict:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFAutoModelWithLMHead.from_pretrained(__lowerCAmelCase   )
    self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )
    self.assertEqual(model.num_parameters()				,				14_410   )
    self.assertEqual(model.num_parameters(only_trainable=__lowerCAmelCase   )				,				14_410   )






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> Dict:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFAutoModelWithLMHead.from_pretrained(__lowerCAmelCase   )
    self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )
    self.assertEqual(model.num_parameters()				,				14_410   )
    self.assertEqual(model.num_parameters(only_trainable=__lowerCAmelCase   )				,				14_410   )






  def 	_UpperCAmelCase				(				self:			Optional[int]   ) -> str:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFAutoModel.from_pretrained("sgugger/funnel-random-tiny"   )
    self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

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

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






  def 	_UpperCAmelCase				(				self:			Tuple   ) -> List[Any]:




    '''simple docstring'''




    try:
      AutoConfig.register("new-model"				,				__lowerCAmelCase   )

      __UpperCAmelCase						    =       [
          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
          __UpperCAmelCase						    =       BertModelTester(self   ).get_config()
          __UpperCAmelCase						    =       NewModelConfig(**tiny_config.to_dict()   )
          __UpperCAmelCase						    =       auto_class.from_config(__lowerCAmelCase   )
          self.assertIsInstance(__lowerCAmelCase				,				__lowerCAmelCase   )

          with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(__lowerCAmelCase   )
            __UpperCAmelCase						    =       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 	_UpperCAmelCase				(				self:			List[Any]   ) -> List[Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Optional[int]   ) -> str:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Any   ) -> List[str]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Dict   ) -> Optional[int]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> str:




    '''simple docstring'''




    __UpperCAmelCase						    =       TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert"   )
    with RequestCounter() as counter:
      __UpperCAmelCase						    =       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
    __UpperCAmelCase						    =       TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded"   )
    with RequestCounter() as counter:
      __UpperCAmelCase						    =       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   )




from ..utils import DummyObject, requires_backends


class 			UpperCAmelCase__    (		metaclass=snake_case					):



  """simple docstring"""






  lowerCAmelCase__		:	List[str]    		= ['transformers', 'torch', 'note_seq']






  def __init__(				self:			List[str]				,				*__lowerCAmelCase:			Optional[int]				,				**__lowerCAmelCase:			List[Any]   ) -> Optional[int]:




    '''simple docstring'''




    requires_backends(self				,				["transformers", "torch", "note_seq"]   )






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Optional[int]				,				*__lowerCAmelCase:			Any				,				**__lowerCAmelCase:			List[str]   ) -> Dict:




    '''simple docstring'''




    requires_backends(cls				,				["transformers", "torch", "note_seq"]   )






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Union[str, Any]				,				*__lowerCAmelCase:			Optional[Any]				,				**__lowerCAmelCase:			Optional[Any]   ) -> Any:




    '''simple docstring'''




    requires_backends(cls				,				["transformers", "torch", "note_seq"]   )



