'''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 lowercase_       (		A__			):







   """simple docstring"""

   __lowerCAmelCase         =   "new-model"


if is_tf_available():







   class lowercase_       (		A__			):







      """simple docstring"""

      __lowerCAmelCase         =   NewModelConfig







@require_tf
class lowercase_       (		unittest.TestCase			):







   """simple docstring"""

   @slow
   def 	__UpperCAmelCase   (      self      :			Any     )  ->			Tuple:
        _A    		=  'bert-base-cased'
        _A    		=  AutoConfig.from_pretrained(_lowerCamelCase     )
        self.assertIsNotNone(_lowerCamelCase     )
        self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )

        _A    		=  TFAutoModel.from_pretrained(_lowerCamelCase     )
        self.assertIsNotNone(_lowerCamelCase     )
        self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			str:
        _A    		=  'bert-base-cased'
        _A    		=  AutoConfig.from_pretrained(_lowerCamelCase     )
        self.assertIsNotNone(_lowerCamelCase     )
        self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )

        _A    		=  TFAutoModelForPreTraining.from_pretrained(_lowerCamelCase     )
        self.assertIsNotNone(_lowerCamelCase     )
        self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Dict:



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

             _A    		=  TFAutoModelForCausalLM.from_pretrained(_lowerCamelCase     )
             _A    		=  TFAutoModelForCausalLM.from_pretrained(_lowerCamelCase, output_loading_info=_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Union[str, Any]:



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

             _A    		=  TFAutoModelWithLMHead.from_pretrained(_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			int     )  ->			List[str]:



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

             _A    		=  TFAutoModelForMaskedLM.from_pretrained(_lowerCamelCase     )
             _A    		=  TFAutoModelForMaskedLM.from_pretrained(_lowerCamelCase, output_loading_info=_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			int:



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

             _A    		=  TFAutoModelForSeqaSeqLM.from_pretrained(_lowerCamelCase     )
             _A    		=  TFAutoModelForSeqaSeqLM.from_pretrained(_lowerCamelCase, output_loading_info=_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Union[str, Any]:



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

             _A    		=  TFAutoModelForSequenceClassification.from_pretrained(_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			List[str]:



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

             _A    		=  TFAutoModelForQuestionAnswering.from_pretrained(_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   @slow
   @require_tensorflow_probability
   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			str:



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

             _A    		=  TFAutoModelForTableQuestionAnswering.from_pretrained(_lowerCamelCase     )
             _A    		=  TFAutoModelForTableQuestionAnswering.from_pretrained(
                 _lowerCamelCase, output_loading_info=_lowerCamelCase     )
             self.assertIsNotNone(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			int:
        _A    		=  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 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			List[str]:
        _A    		=  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 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			str:
        _A    		=  TFAutoModel.from_pretrained('sgugger/funnel-random-tiny'     )
        self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )

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

        with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(_lowerCamelCase     )
             _A    		=  TFAutoModel.from_pretrained(_lowerCamelCase     )
             self.assertIsInstance(_lowerCamelCase, _lowerCamelCase     )
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Union[str, Any]:



        try:
             AutoConfig.register('new-model', _lowerCamelCase     )

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

                       with tempfile.TemporaryDirectory() as tmp_dir:
                            model.save_pretrained(_lowerCamelCase     )
                            _A    		=  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      :			Dict     )  ->			str:



        with self.assertRaisesRegex(
            _lowerCamelCase, 'bert-base is not a local folder and is not a valid model identifier'     ):
             _A    		=  TFAutoModel.from_pretrained('bert-base'     )
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			Tuple:



        with self.assertRaisesRegex(
            _lowerCamelCase, r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'     ):
             _A    		=  TFAutoModel.from_pretrained(_lowerCamelCase, revision='aaaaaa'     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Optional[int]:



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



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

   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Dict:
        _A    		=  TFAutoModel.from_pretrained('hf-internal-testing/tiny-random-bert'     )
        with RequestCounter() as counter:
             _A    		=  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
        _A    		=  TFAutoModel.from_pretrained('ArthurZ/tiny-random-bert-sharded'     )
        with RequestCounter() as counter:
             _A    		=  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 math

import flax.linen as nn
import jax.numpy as jnp




def       A_  (				__lowercase						,	__lowercase						,	__lowercase = 1						,	__lowercase = 1						,	__lowercase = 1.0e4						,	__lowercase = False						,	__lowercase = 1.0						,	):
 assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
 assert embedding_dim % 2 == 0, F'''Embedding dimension {embedding_dim} should be even'''
 UpperCamelCase_     :					Optional[int]   	=float(embedding_dim // 2	)
 UpperCamelCase_     :					Optional[Any]   	=math.log(max_timescale / min_timescale	) / (num_timescales - freq_shift)
 UpperCamelCase_     :					List[Any]   	=min_timescale * jnp.exp(jnp.arange(__lowercase						,	dtype=jnp.floataa	) * -log_timescale_increment	)
 UpperCamelCase_     :					int   	=jnp.expand_dims(__lowercase						,	1	) * jnp.expand_dims(__lowercase						,	0	)

 # scale embeddings
 UpperCamelCase_     :					List[str]   	=scale * emb

 if flip_sin_to_cos:
  UpperCamelCase_     :					Tuple   	=jnp.concatenate([jnp.cos(__lowercase	), jnp.sin(__lowercase	)]						,	axis=1	)
 else:
  UpperCamelCase_     :					Tuple   	=jnp.concatenate([jnp.sin(__lowercase	), jnp.cos(__lowercase	)]						,	axis=1	)
 UpperCamelCase_     :					List[Any]   	=jnp.reshape(__lowercase						,	[jnp.shape(__lowercase	)[0], embedding_dim]	)
 return signal





class 				a__					(						nn.Module       ):
     UpperCAmelCase__         =      32
     UpperCAmelCase__         =      jnp.floataa



     @nn.compact
     def __call__(							self :Optional[Any]      ,  _lowerCamelCase :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Dict   	=nn.Dense(self.time_embed_dim      ,  dtype=self.dtype      ,  name='linear_1'   )(_lowerCamelCase   )
      UpperCamelCase_     :					Any   	=nn.silu(_lowerCamelCase   )
      UpperCamelCase_     :					Tuple   	=nn.Dense(self.time_embed_dim      ,  dtype=self.dtype      ,  name='linear_2'   )(_lowerCamelCase   )
      return temb





class 				a__					(						nn.Module       ):
     UpperCAmelCase__         =      32
     UpperCAmelCase__         =      False
     UpperCAmelCase__         =      1



     @nn.compact
     def __call__(							self :Union[str, Any]      ,  _lowerCamelCase :Dict   ):







      '''simple docstring'''



      return get_sinusoidal_embeddings(
          _lowerCamelCase      ,  embedding_dim=self.dim      ,  flip_sin_to_cos=self.flip_sin_to_cos      ,  freq_shift=self.freq_shift   )
