





import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
	from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments



@require_torch
class 					__lowerCAmelCase       (			unittest.TestCase     ):






			def 			lowerCamelCase  (self					,							__magic_name__    )     ->   List[Any]:


							'''simple docstring'''



							for model_result in results.values():
											for batch_size, sequence_length in zip(model_result['''bs''']					,							model_result['''ss''']    ):
															snake_case_						:    Dict				=  model_result['''result'''][batch_size][sequence_length]
															self.assertIsNotNone(__magic_name__    )






			def 			lowerCamelCase  (self    )     ->   str:


							'''simple docstring'''



							snake_case_						:    str				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    List[Any]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    Any				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    Tuple				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   Dict:


							'''simple docstring'''



							snake_case_						:    Tuple				=  '''sgugger/tiny-distilbert-classification'''
							snake_case_						:    Optional[Any]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							only_pretrain_model=__magic_name__					,							)
							snake_case_						:    List[str]				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    Union[str, Any]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							snake_case_						:    int				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Union[str, Any]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							torchscript=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    int				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    List[str]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			@unittest.skipIf(torch_device == '''cpu'''					,							'''Cant do half precision'''    )
			def 			lowerCamelCase  (self    )     ->   Dict:


							'''simple docstring'''



							snake_case_						:    Optional[int]				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Optional[int]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							fpaa=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    int				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    Optional[Any]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   List[Any]:


							'''simple docstring'''



							snake_case_						:    str				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Any				=  AutoConfig.from_pretrained(__magic_name__    )
							# set architectures equal to `None`
							snake_case_						:    Optional[Any]				=  None
							snake_case_						:    str				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    List[str]				=  PyTorchBenchmark(__magic_name__					,							configs=[config]    )
							snake_case_						:    Union[str, Any]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							snake_case_						:    List[str]				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    str				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    List[Any]				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    int				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_train_result    )
							self.check_results_dict_not_empty(results.memory_train_result    )






			@unittest.skipIf(torch_device == '''cpu'''					,							'''Can\'t do half precision'''    )
			def 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							snake_case_						:    Union[str, Any]				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Optional[int]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							fpaa=__magic_name__					,							multi_process=__magic_name__					,							)
							snake_case_						:    List[Any]				=  PyTorchBenchmark(__magic_name__    )
							snake_case_						:    Optional[int]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_train_result    )
							self.check_results_dict_not_empty(results.memory_train_result    )






			def 			lowerCamelCase  (self    )     ->   Dict:


							'''simple docstring'''



							snake_case_						:    Union[str, Any]				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Union[str, Any]				=  AutoConfig.from_pretrained(__magic_name__    )
							snake_case_						:    Optional[int]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    Optional[int]				=  PyTorchBenchmark(__magic_name__					,							configs=[config]    )
							snake_case_						:    Optional[int]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   Any:


							'''simple docstring'''



							snake_case_						:    Union[str, Any]				=  '''sshleifer/tinier_bart'''
							snake_case_						:    Any				=  AutoConfig.from_pretrained(__magic_name__    )
							snake_case_						:    Tuple				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    Tuple				=  PyTorchBenchmark(__magic_name__					,							configs=[config]    )
							snake_case_						:    Tuple				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_inference_result    )
							self.check_results_dict_not_empty(results.memory_inference_result    )






			def 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							snake_case_						:    str				=  '''sshleifer/tiny-gpt2'''
							snake_case_						:    Dict				=  AutoConfig.from_pretrained(__magic_name__    )
							snake_case_						:    List[str]				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    Union[str, Any]				=  PyTorchBenchmark(__magic_name__					,							configs=[config]    )
							snake_case_						:    Optional[int]				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_train_result    )
							self.check_results_dict_not_empty(results.memory_train_result    )






			def 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							snake_case_						:    Optional[int]				=  '''sshleifer/tinier_bart'''
							snake_case_						:    Union[str, Any]				=  AutoConfig.from_pretrained(__magic_name__    )
							snake_case_						:    str				=  PyTorchBenchmarkArguments(
							    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							multi_process=__magic_name__					,							)
							snake_case_						:    List[Any]				=  PyTorchBenchmark(__magic_name__					,							configs=[config]    )
							snake_case_						:    str				=  benchmark.run()
							self.check_results_dict_not_empty(results.time_train_result    )
							self.check_results_dict_not_empty(results.memory_train_result    )






			def 			lowerCamelCase  (self    )     ->   List[str]:


							'''simple docstring'''



							snake_case_						:    int				=  '''sshleifer/tiny-gpt2'''
							with tempfile.TemporaryDirectory() as tmp_dir:
											snake_case_						:    Optional[int]				=  PyTorchBenchmarkArguments(
											    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							save_to_csv=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							inference_time_csv_file=os.path.join(__magic_name__					,							'''inf_time.csv'''    )					,							train_memory_csv_file=os.path.join(__magic_name__					,							'''train_mem.csv'''    )					,							inference_memory_csv_file=os.path.join(__magic_name__					,							'''inf_mem.csv'''    )					,							train_time_csv_file=os.path.join(__magic_name__					,							'''train_time.csv'''    )					,							env_info_csv_file=os.path.join(__magic_name__					,							'''env.csv'''    )					,							multi_process=__magic_name__					,							)
											snake_case_						:    List[str]				=  PyTorchBenchmark(__magic_name__    )
											benchmark.run()
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''inf_time.csv'''    )    ).exists()    )
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''train_time.csv'''    )    ).exists()    )
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''inf_mem.csv'''    )    ).exists()    )
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''train_mem.csv'''    )    ).exists()    )
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''env.csv'''    )    ).exists()    )






			def 			lowerCamelCase  (self    )     ->   List[str]:


							'''simple docstring'''



							snake_case_						:    str				=  '''sshleifer/tiny-gpt2'''

							def _check_summary_is_not_empty(__magic_name__    ):
											self.assertTrue(hasattr(__magic_name__					,							'''sequential'''    )    )
											self.assertTrue(hasattr(__magic_name__					,							'''cumulative'''    )    )
											self.assertTrue(hasattr(__magic_name__					,							'''current'''    )    )
											self.assertTrue(hasattr(__magic_name__					,							'''total'''    )    )

							with tempfile.TemporaryDirectory() as tmp_dir:
											snake_case_						:    Dict				=  PyTorchBenchmarkArguments(
											    models=[MODEL_ID]					,							training=__magic_name__					,							inference=__magic_name__					,							sequence_lengths=[8]					,							batch_sizes=[1]					,							log_filename=os.path.join(__magic_name__					,							'''log.txt'''    )					,							log_print=__magic_name__					,							trace_memory_line_by_line=__magic_name__					,							multi_process=__magic_name__					,							)
											snake_case_						:    Union[str, Any]				=  PyTorchBenchmark(__magic_name__    )
											snake_case_						:    Union[str, Any]				=  benchmark.run()
											_check_summary_is_not_empty(result.inference_summary    )
											_check_summary_is_not_empty(result.train_summary    )
											self.assertTrue(Path(os.path.join(__magic_name__					,							'''log.txt'''    )    ).exists()    )







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


lowerCAmelCase_    =				logging.get_logger(__name__)

lowerCAmelCase_    =				{
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}



class 					__lowerCAmelCase       (			_a     ):
			lowerCamelCase_     :     Any       			=      '''megatron-bert'''






			def __init__(self					,							__magic_name__=2_9056					,							__magic_name__=1024					,							__magic_name__=24					,							__magic_name__=16					,							__magic_name__=4096					,							__magic_name__="gelu"					,							__magic_name__=0.1					,							__magic_name__=0.1					,							__magic_name__=512					,							__magic_name__=2					,							__magic_name__=0.02					,							__magic_name__=1e-12					,							__magic_name__=0					,							__magic_name__="absolute"					,							__magic_name__=True					,							**__magic_name__					,							)     ->   Any:


							'''simple docstring'''



							super().__init__(pad_token_id=__magic_name__					,							**__magic_name__    )

							snake_case_						:    Union[str, Any]				=  vocab_size
							snake_case_						:    Dict				=  hidden_size
							snake_case_						:    Dict				=  num_hidden_layers
							snake_case_						:    Optional[int]				=  num_attention_heads
							snake_case_						:    int				=  hidden_act
							snake_case_						:    List[str]				=  intermediate_size
							snake_case_						:    Dict				=  hidden_dropout_prob
							snake_case_						:    str				=  attention_probs_dropout_prob
							snake_case_						:    str				=  max_position_embeddings
							snake_case_						:    Any				=  type_vocab_size
							snake_case_						:    int				=  initializer_range
							snake_case_						:    int				=  layer_norm_eps
							snake_case_						:    List[str]				=  position_embedding_type
							snake_case_						:    Dict				=  use_cache
