


'''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 		_snake_case					(     a__	):
					lowerCAmelCase      :List[str]   =					'''new-model'''


if is_tf_available():


		class 		_snake_case					(     a__	):
							lowerCAmelCase      :List[str]   =					NewModelConfig


@require_tf
class 		_snake_case					(     unittest.TestCase	):




					@slow
					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Any        =     """bert-base-cased"""
								UpperCAmelCase__  :					str        =     AutoConfig.from_pretrained(_lowerCamelCase)
								self.assertIsNotNone(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					Union[str, Any]        =     TFAutoModel.from_pretrained(_lowerCamelCase)
								self.assertIsNotNone(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)




					@slow
					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[Any]        =     """bert-base-cased"""
								UpperCAmelCase__  :					Optional[Any]        =     AutoConfig.from_pretrained(_lowerCamelCase)
								self.assertIsNotNone(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					Optional[Any]        =     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]:
											UpperCAmelCase__  :					Optional[int]        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					Union[str, Any]        =     TFAutoModelForCausalLM.from_pretrained(_lowerCamelCase)
											UpperCAmelCase__   ,       UpperCAmelCase__  :					str        =     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]:
											UpperCAmelCase__  :					int        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					List[Any]        =     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]:
											UpperCAmelCase__  :					Tuple        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					str        =     TFAutoModelForMaskedLM.from_pretrained(_lowerCamelCase)
											UpperCAmelCase__   ,       UpperCAmelCase__  :					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]:
											UpperCAmelCase__  :					int        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					Tuple        =     TFAutoModelForSeqaSeqLM.from_pretrained(_lowerCamelCase)
											UpperCAmelCase__   ,       UpperCAmelCase__  :					Tuple        =     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"]:
											UpperCAmelCase__  :					Tuple        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					List[str]        =     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"]:
											UpperCAmelCase__  :					Optional[int]        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					Optional[int]        =     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]:
											UpperCAmelCase__  :					Dict        =     AutoConfig.from_pretrained(_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

											UpperCAmelCase__  :					str        =     TFAutoModelForTableQuestionAnswering.from_pretrained(_lowerCamelCase)
											UpperCAmelCase__   ,       UpperCAmelCase__  :					Optional[Any]        =     TFAutoModelForTableQuestionAnswering.from_pretrained(
											    _lowerCamelCase							,							output_loading_info=_lowerCamelCase)
											self.assertIsNotNone(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[Any]        =     TFAutoModelWithLMHead.from_pretrained(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)
								self.assertEqual(model.num_parameters()							,							1_4410)
								self.assertEqual(model.num_parameters(only_trainable=_lowerCamelCase)							,							1_4410)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					str        =     TFAutoModelWithLMHead.from_pretrained(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)
								self.assertEqual(model.num_parameters()							,							1_4410)
								self.assertEqual(model.num_parameters(only_trainable=_lowerCamelCase)							,							1_4410)




					def 				snake_case__						(      self):
								# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
								UpperCAmelCase__  :					int        =     TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny""")
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					Tuple        =     copy.deepcopy(model.config)
								UpperCAmelCase__  :					List[str]        =     ["""FunnelBaseModel"""]
								UpperCAmelCase__  :					Optional[int]        =     TFAutoModel.from_config(_lowerCamelCase)
								self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

								with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(_lowerCamelCase)
											UpperCAmelCase__  :					Any        =     TFAutoModel.from_pretrained(_lowerCamelCase)
											self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)




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

											UpperCAmelCase__  :					Any        =     [
											    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__  :					Dict        =     BertModelTester(self).get_config()
																	UpperCAmelCase__  :					Optional[int]        =     NewModelConfig(**tiny_config.to_dict())
																	UpperCAmelCase__  :					List[Any]        =     auto_class.from_config(_lowerCamelCase)
																	self.assertIsInstance(_lowerCamelCase							,							_lowerCamelCase)

																	with tempfile.TemporaryDirectory() as tmp_dir:
																				model.save_pretrained(_lowerCamelCase)
																				UpperCAmelCase__  :					Dict        =     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"""):
											UpperCAmelCase__  :					str        =     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\)"""):
											UpperCAmelCase__  :					Optional[Any]        =     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"""							,							):
											UpperCAmelCase__  :					Dict        =     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"""):
											UpperCAmelCase__  :					int        =     TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only""")




					def 				snake_case__						(      self):
								# Make sure we have cached the model.
								UpperCAmelCase__  :					Tuple        =     TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert""")
								with RequestCounter() as counter:
											UpperCAmelCase__  :					Dict        =     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__  :					int        =     TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded""")
								with RequestCounter() as counter:
											UpperCAmelCase__  :					Any        =     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 importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


__A	='scheduler_config.json'


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Union[str, Any]   =					1
					lowerCAmelCase      :Union[str, Any]   =					2
					lowerCAmelCase      :Optional[int]   =					3
					lowerCAmelCase      :Optional[int]   =					4
					lowerCAmelCase      :Any   =					5
					lowerCAmelCase      :Tuple   =					6
					lowerCAmelCase      :List[Any]   =					7
					lowerCAmelCase      :str   =					8
					lowerCAmelCase      :List[Any]   =					9
					lowerCAmelCase      :List[str]   =					10
					lowerCAmelCase      :Union[str, Any]   =					11
					lowerCAmelCase      :Optional[int]   =					12
					lowerCAmelCase      :str   =					13
					lowerCAmelCase      :Dict   =					14


@dataclass
class 		_snake_case					(     a__	):
					lowerCAmelCase      :torch.FloatTensor


class 		_snake_case					:
					lowerCAmelCase      :str   =					SCHEDULER_CONFIG_NAME
					lowerCAmelCase      :Union[str, Any]   =					[]
					lowerCAmelCase      :List[str]   =					True




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							_lowerCamelCase=False							,							**_lowerCamelCase							,							):
								UpperCAmelCase__   ,       UpperCAmelCase__   ,       UpperCAmelCase__  :					Optional[int]        =     cls.load_config(
								    pretrained_model_name_or_path=_lowerCamelCase							,							subfolder=_lowerCamelCase							,							return_unused_kwargs=_lowerCamelCase							,							return_commit_hash=_lowerCamelCase							,							**_lowerCamelCase							,							)
								return cls.from_config(_lowerCamelCase							,							return_unused_kwargs=_lowerCamelCase							,							**_lowerCamelCase)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = False							,							**_lowerCamelCase):
								self.save_config(save_directory=_lowerCamelCase							,							push_to_hub=_lowerCamelCase							,							**_lowerCamelCase)




					@property
					def 				snake_case__						(      self):
								return self._get_compatibles()




					@classmethod
					def 				snake_case__						(      cls):
								UpperCAmelCase__  :					List[Any]        =     list(set([cls.__name__] + cls._compatibles))
								UpperCAmelCase__  :					int        =     importlib.import_module(__name__.split(""".""")[0])
								UpperCAmelCase__  :					Dict        =     [
								    getattr(_lowerCamelCase							,							_lowerCamelCase) for c in compatible_classes_str if hasattr(_lowerCamelCase							,							_lowerCamelCase)
								]
								return compatible_classes