





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    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''
 lowercase_		=			"new-model"


if is_tf_available():



     class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







      '''simple docstring'''
      lowercase_		=			NewModelConfig



@require_tf
class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ):







 '''simple docstring'''

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



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  ="bert-base-cased"
   lowerCamelCase__:      Tuple			  =AutoConfig.from_pretrained(UpperCAmelCase_)
   self.assertIsNotNone(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

   lowerCamelCase__:      Any			  =TFAutoModel.from_pretrained(UpperCAmelCase_)
   self.assertIsNotNone(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:      Any			  ="bert-base-cased"
   lowerCamelCase__:      Tuple			  =AutoConfig.from_pretrained(UpperCAmelCase_)
   self.assertIsNotNone(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

   lowerCamelCase__:      Tuple			  =TFAutoModelForPreTraining.from_pretrained(UpperCAmelCase_)
   self.assertIsNotNone(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



   for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      str			  =AutoConfig.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

     lowerCamelCase__:      List[str]			  =TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase_)
     lowerCamelCase__  ,					lowerCamelCase__:      Optional[int]			  =TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase_       , output_loading_info=UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



   for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      List[Any]			  =AutoConfig.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

     lowerCamelCase__:      Any			  =TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->int:



   '''simple docstring'''



   for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      Optional[int]			  =AutoConfig.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

     lowerCamelCase__:      Dict			  =TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase_)
     lowerCamelCase__  ,					lowerCamelCase__:      List[str]			  =TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase_       , output_loading_info=UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



   for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      str			  =AutoConfig.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

     lowerCamelCase__:      str			  =TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_)
     lowerCamelCase__  ,					lowerCamelCase__:      List[str]			  =TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_       , output_loading_info=UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Any:



   '''simple docstring'''



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

     lowerCamelCase__:      Dict			  =TFAutoModelForSequenceClassification.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

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



   '''simple docstring'''



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

     lowerCamelCase__:      Tuple			  =TFAutoModelForQuestionAnswering.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 @slow
 @require_tensorflow_probability
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->List[str]:



   '''simple docstring'''



   for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
     lowerCamelCase__:      Tuple			  =AutoConfig.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

     lowerCamelCase__:      Optional[Any]			  =TFAutoModelForTableQuestionAnswering.from_pretrained(UpperCAmelCase_)
     lowerCamelCase__  ,					lowerCamelCase__:      List[str]			  =TFAutoModelForTableQuestionAnswering.from_pretrained(
         UpperCAmelCase_       , output_loading_info=UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->int:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)
   self.assertEqual(model.num_parameters()       , 14_410)
   self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase_)       , 14_410)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)
   self.assertEqual(model.num_parameters()       , 14_410)
   self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase_)       , 14_410)

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



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =TFAutoModel.from_pretrained("sgugger/funnel-random-tiny")
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

   lowerCamelCase__:      str			  =copy.deepcopy(model.config)
   lowerCamelCase__:      Dict			  =["FunnelBaseModel"]
   lowerCamelCase__:      List[Any]			  =TFAutoModel.from_config(UpperCAmelCase_)
   self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

   with tempfile.TemporaryDirectory() as tmp_dir:
     model.save_pretrained(UpperCAmelCase_)
     lowerCamelCase__:      List[Any]			  =TFAutoModel.from_pretrained(UpperCAmelCase_)
     self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->int:



   '''simple docstring'''



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

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

         # Now that the config is registered, it can be used as any other config with the auto-API
         lowerCamelCase__:      Tuple			  =BertModelTester(self).get_config()
         lowerCamelCase__:      List[str]			  =NewModelConfig(**tiny_config.to_dict())
         lowerCamelCase__:      List[Any]			  =auto_class.from_config(UpperCAmelCase_)
         self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

         with tempfile.TemporaryDirectory() as tmp_dir:
           model.save_pretrained(UpperCAmelCase_)
           lowerCamelCase__:      Tuple			  =auto_class.from_pretrained(UpperCAmelCase_)
           self.assertIsInstance(UpperCAmelCase_       , UpperCAmelCase_)

   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 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->int:



   '''simple docstring'''



   with self.assertRaisesRegex(
       UpperCAmelCase_       , "bert-base is not a local folder and is not a valid model identifier"):
     lowerCamelCase__:      Optional[Any]			  =TFAutoModel.from_pretrained("bert-base")

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



   '''simple docstring'''



   with self.assertRaisesRegex(
       UpperCAmelCase_       , R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"):
     lowerCamelCase__:      Optional[int]			  =TFAutoModel.from_pretrained(UpperCAmelCase_       , revision="aaaaaa")

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



   '''simple docstring'''



   with self.assertRaisesRegex(
       UpperCAmelCase_       , "hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"       , ):
     lowerCamelCase__:      Union[str, Any]			  =TFAutoModel.from_pretrained("hf-internal-testing/config-no-model")

 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Dict:



   '''simple docstring'''



   with self.assertRaisesRegex(UpperCAmelCase_       , "Use `from_pt=True` to load this model"):
     lowerCamelCase__:      List[Any]			  =TFAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only")







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



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =TFAutoModel.from_pretrained("hf-internal-testing/tiny-random-bert")
   with RequestCounter() as counter:
     lowerCamelCase__:      List[str]			  =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__:      Optional[int]			  =TFAutoModel.from_pretrained("ArthurZ/tiny-random-bert-sharded")
   with RequestCounter() as counter:
     lowerCamelCase__:      List[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 gc
import threading
import time

import psutil
import torch



class 			A__   :




	'''simple docstring'''





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

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[Any]												=       psutil.Process()
				__lowerCAmelCase      :  str												=       False

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[Any]												=       -1

				while True:
							__lowerCAmelCase      :  str												=       max(self.process.memory_info().rss       ,			self.cpu_memory_peak)

							# can't sleep or will not catch the peak right (this comment is here on purpose)
							if not self.peak_monitoring:
										break

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple)				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       True
				__lowerCAmelCase      :  str												=       threading.Thread(target=self.peak_monitor)
				__lowerCAmelCase      :  Tuple												=       True
				self.thread.start()







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

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       False
				self.thread.join()
				return self.cpu_memory_peak






__snake_case							:   Tuple      =    PeakCPUMemory()


def 			_lowercase				(  )						->   str:
			# Time
			__lowerCAmelCase      :  str												=       {"time": time.time()}

			gc.collect()
			torch.cuda.empty_cache()

			# CPU mem
			__lowerCAmelCase      :  Optional[Any]												=       psutil.Process().memory_info().rss
			cpu_peak_tracker.start()

			# GPU mem
			for i in range(torch.cuda.device_count()	):
						__lowerCAmelCase      :  Union[str, Any]												=       torch.cuda.memory_allocated(__snake_case	)
			torch.cuda.reset_peak_memory_stats()

			return measures


def 			_lowercase				(  __snake_case	)						->   Optional[Any]:
			# Time
			__lowerCAmelCase      :  str												=       {"time": time.time() - start_measures["time"]}

			gc.collect()
			torch.cuda.empty_cache()

			# CPU mem
			__lowerCAmelCase      :  str												=       (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
			__lowerCAmelCase      :  List[str]												=       (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20

			# GPU mem
			for i in range(torch.cuda.device_count()	):
						__lowerCAmelCase      :  Union[str, Any]												=       (torch.cuda.memory_allocated(__snake_case	) - start_measures[str(__snake_case	)]) / 2**20
						__lowerCAmelCase      :  Any												=       (torch.cuda.max_memory_allocated(__snake_case	) - start_measures[str(__snake_case	)]) / 2**20

			return measures





def 			_lowercase				(  __snake_case					,__snake_case	)						->   Dict:
			print(F"""{description}:"""	)
			print(F"""- Time: {measures['time']:.2f}s"""	)
			for i in range(torch.cuda.device_count()	):
						print(F"""- GPU {i} allocated: {measures[str(__snake_case	)]:.2f}MiB"""	)
						__lowerCAmelCase      :  Optional[Any]												=       measures[F"""{i}-peak"""]
						print(F"""- GPU {i} peak: {peak:.2f}MiB"""	)
			print(F"""- CPU RAM allocated: {measures['cpu']:.2f}MiB"""	)
			print(F"""- CPU RAM peak: {measures['cpu-peak']:.2f}MiB"""	)