





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    (    _UpperCamelCase    ):

  """simple docstring"""


  a_        = """new-model"""


if is_tf_available():



     class 				_UpperCAmelCase    (    _UpperCamelCase    ):

       """simple docstring"""


       a_        = NewModelConfig



@require_tf
class 				_UpperCAmelCase    (    unittest.TestCase    ):

  """simple docstring"""


  @slow
  def 				lowercase		(				self :		Dict     )							->						Optional[Any]:
     __lowerCAmelCase								=	'bert-base-cased'
     __lowerCAmelCase								=	AutoConfig.from_pretrained(lowerCAmelCase_     )
     self.assertIsNotNone(lowerCAmelCase_     )
     self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )

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






  @slow
  def 				lowercase		(				self :		List[Any]     )							->						List[Any]:
     __lowerCAmelCase								=	'bert-base-cased'
     __lowerCAmelCase								=	AutoConfig.from_pretrained(lowerCAmelCase_     )
     self.assertIsNotNone(lowerCAmelCase_     )
     self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )

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






  @slow
  def 				lowercase		(				self :		Tuple     )							->						int:
     for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase								=	AutoConfig.from_pretrained(lowerCAmelCase_     )
        self.assertIsNotNone(lowerCAmelCase_     )
        self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )

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






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

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






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

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






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

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






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

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






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

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






  @slow
  @require_tensorflow_probability
  def 				lowercase		(				self :		str     )							->						List[str]:
     for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
        __lowerCAmelCase								=	AutoConfig.from_pretrained(lowerCAmelCase_     )
        self.assertIsNotNone(lowerCAmelCase_     )
        self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )

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






  def 				lowercase		(				self :		int     )							->						List[Any]:
     __lowerCAmelCase								=	TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase_     )
     self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )
     self.assertEqual(model.num_parameters()					,  1_4_4_1_0     )
     self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase_     )					,  1_4_4_1_0     )






  def 				lowercase		(				self :		Optional[int]     )							->						List[Any]:
     __lowerCAmelCase								=	TFAutoModelWithLMHead.from_pretrained(lowerCAmelCase_     )
     self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )
     self.assertEqual(model.num_parameters()					,  1_4_4_1_0     )
     self.assertEqual(model.num_parameters(only_trainable=lowerCAmelCase_     )					,  1_4_4_1_0     )






  def 				lowercase		(				self :		Dict     )							->						List[str]:
     # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
     __lowerCAmelCase								=	TFAutoModel.from_pretrained('sgugger/funnel-random-tiny'     )
     self.assertIsInstance(lowerCAmelCase_					,  lowerCAmelCase_     )

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

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






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

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

              with tempfile.TemporaryDirectory() as tmp_dir:
                 model.save_pretrained(lowerCAmelCase_     )
                 __lowerCAmelCase								=	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 				lowercase		(				self :		List[str]     )							->						Union[str, Any]:
     with self.assertRaisesRegex(
         lowerCAmelCase_					,  'bert-base is not a local folder and is not a valid model identifier'     ):
        __lowerCAmelCase								=	TFAutoModel.from_pretrained('bert-base'     )






  def 				lowercase		(				self :		int     )							->						str:
     with self.assertRaisesRegex(
         lowerCAmelCase_					,  R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'     ):
        __lowerCAmelCase								=	TFAutoModel.from_pretrained(lowerCAmelCase_					,  revision='aaaaaa'     )






  def 				lowercase		(				self :		int     )							->						int:
     with self.assertRaisesRegex(
         lowerCAmelCase_					,  '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 				lowercase		(				self :		Union[str, Any]     )							->						Optional[int]:
     with self.assertRaisesRegex(lowerCAmelCase_					,  'Use `from_pt=True` to load this model'     ):
        __lowerCAmelCase								=	TFAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'     )






  def 				lowercase		(				self :		Optional[int]     )							->						str:
     # 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     )








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_snake_case  :    List[Any]						    =  logging.get_logger(__name__)

_snake_case  :    Any						    =  {
    'google/pegasus-large': 'https://huggingface.co/google/pegasus-large/resolve/main/config.json',
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}



class 				_UpperCAmelCase    (    _UpperCamelCase    ):

  """simple docstring"""


  a_        = """pegasus"""
  a_        = ["""past_key_values"""]
  a_        = {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






  def __init__(				self :		List[str]					,  lowerCAmelCase_ :		Union[str, Any]=5_0_2_6_5					,  lowerCAmelCase_ :		Union[str, Any]=1_0_2_4					,  lowerCAmelCase_ :		Union[str, Any]=1_2					,  lowerCAmelCase_ :		Dict=4_0_9_6					,  lowerCAmelCase_ :		str=1_6					,  lowerCAmelCase_ :		List[Any]=1_2					,  lowerCAmelCase_ :		Union[str, Any]=4_0_9_6					,  lowerCAmelCase_ :		Union[str, Any]=1_6					,  lowerCAmelCase_ :		Any=0.0					,  lowerCAmelCase_ :		Optional[int]=0.0					,  lowerCAmelCase_ :		List[Any]=True					,  lowerCAmelCase_ :		Any=True					,  lowerCAmelCase_ :		Dict="gelu"					,  lowerCAmelCase_ :		Optional[int]=1_0_2_4					,  lowerCAmelCase_ :		List[str]=0.1					,  lowerCAmelCase_ :		str=0.0					,  lowerCAmelCase_ :		Union[str, Any]=0.0					,  lowerCAmelCase_ :		Optional[int]=0.02					,  lowerCAmelCase_ :		Tuple=0					,  lowerCAmelCase_ :		Tuple=False					,  lowerCAmelCase_ :		Union[str, Any]=0					,  lowerCAmelCase_ :		int=1					,  lowerCAmelCase_ :		Tuple=1					,  **lowerCAmelCase_ :		Tuple					,  )							->						List[str]:
     __lowerCAmelCase								=	vocab_size
     __lowerCAmelCase								=	max_position_embeddings
     __lowerCAmelCase								=	d_model
     __lowerCAmelCase								=	encoder_ffn_dim
     __lowerCAmelCase								=	encoder_layers
     __lowerCAmelCase								=	encoder_attention_heads
     __lowerCAmelCase								=	decoder_ffn_dim
     __lowerCAmelCase								=	decoder_layers
     __lowerCAmelCase								=	decoder_attention_heads
     __lowerCAmelCase								=	dropout
     __lowerCAmelCase								=	attention_dropout
     __lowerCAmelCase								=	activation_dropout
     __lowerCAmelCase								=	activation_function
     __lowerCAmelCase								=	init_std
     __lowerCAmelCase								=	encoder_layerdrop
     __lowerCAmelCase								=	decoder_layerdrop
     __lowerCAmelCase								=	use_cache
     __lowerCAmelCase								=	encoder_layers
     __lowerCAmelCase								=	scale_embedding  # scale factor will be sqrt(d_model) if True
     super().__init__(
         pad_token_id=lowerCAmelCase_					,  eos_token_id=lowerCAmelCase_					,  is_encoder_decoder=lowerCAmelCase_					,  decoder_start_token_id=lowerCAmelCase_					,  forced_eos_token_id=lowerCAmelCase_					,  **lowerCAmelCase_					,  )






  @property
  def 				lowercase		(				self :		List[Any]     )							->						int:
     return self.encoder_attention_heads






  @property
  def 				lowercase		(				self :		Optional[Any]     )							->						int:
     return self.d_model

