




import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
					from transformers.models.esm.modeling_esm import (
					    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
					    EsmEmbeddings,
					    create_position_ids_from_input_ids,
					)
class        __snake_case  :



				def __init__(      self      :		str		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Optional[int]=13		,	_UpperCAmelCase      :		int=7		,	_UpperCAmelCase      :		Optional[Any]=False		,	_UpperCAmelCase      :		Dict=True		,	_UpperCAmelCase      :		Union[str, Any]=False		,	_UpperCAmelCase      :		Optional[int]=True		,	_UpperCAmelCase      :		Optional[int]=33		,	_UpperCAmelCase      :		int=32		,	_UpperCAmelCase      :		List[Any]=5		,	_UpperCAmelCase      :		Optional[Any]=4		,	_UpperCAmelCase      :		str=37		,	_UpperCAmelCase      :		Any="gelu"		,	_UpperCAmelCase      :		List[Any]=0.1		,	_UpperCAmelCase      :		Tuple=0.1		,	_UpperCAmelCase      :		Any=512		,	_UpperCAmelCase      :		List[str]=16		,	_UpperCAmelCase      :		Union[str, Any]=2		,	_UpperCAmelCase      :		Optional[int]=0.02		,	_UpperCAmelCase      :		Optional[int]=3		,	_UpperCAmelCase      :		Union[str, Any]=4		,	_UpperCAmelCase      :		List[Any]=None		,	)  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      parent
								_lowerCAmelCase    :    Union[str, Any]										=      batch_size
								_lowerCAmelCase    :    List[Any]										=      seq_length
								_lowerCAmelCase    :    str										=      is_training
								_lowerCAmelCase    :    Tuple										=      use_input_mask
								_lowerCAmelCase    :    Dict										=      use_token_type_ids
								_lowerCAmelCase    :    List[Any]										=      use_labels
								_lowerCAmelCase    :    int										=      vocab_size
								_lowerCAmelCase    :    Tuple										=      hidden_size
								_lowerCAmelCase    :    List[str]										=      num_hidden_layers
								_lowerCAmelCase    :    List[Any]										=      num_attention_heads
								_lowerCAmelCase    :    Dict										=      intermediate_size
								_lowerCAmelCase    :    str										=      hidden_act
								_lowerCAmelCase    :    Optional[int]										=      hidden_dropout_prob
								_lowerCAmelCase    :    Dict										=      attention_probs_dropout_prob
								_lowerCAmelCase    :    Union[str, Any]										=      max_position_embeddings
								_lowerCAmelCase    :    Union[str, Any]										=      type_vocab_size
								_lowerCAmelCase    :    List[Any]										=      type_sequence_label_size
								_lowerCAmelCase    :    List[str]										=      initializer_range
								_lowerCAmelCase    :    str										=      num_labels
								_lowerCAmelCase    :    Optional[int]										=      num_choices
								_lowerCAmelCase    :    int										=      scope



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]      )  ->    Dict:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size      )

								_lowerCAmelCase    :    str										=      None
								if self.use_input_mask:
												_lowerCAmelCase    :    str										=      random_attention_mask([self.batch_size, self.seq_length]      )

								_lowerCAmelCase    :    Optional[int]										=      None
								_lowerCAmelCase    :    Union[str, Any]										=      None
								_lowerCAmelCase    :    List[Any]										=      None
								if self.use_labels:
												_lowerCAmelCase    :    Optional[int]										=      ids_tensor([self.batch_size]		,	self.type_sequence_label_size      )
												_lowerCAmelCase    :    str										=      ids_tensor([self.batch_size, self.seq_length]		,	self.num_labels      )
												_lowerCAmelCase    :    List[Any]										=      ids_tensor([self.batch_size]		,	self.num_choices      )

								_lowerCAmelCase    :    Union[str, Any]										=      self.get_config()

								return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



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





								'''simple docstring'''




								return EsmConfig(
								    vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	pad_token_id=1		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	intermediate_size=self.intermediate_size		,	hidden_act=self.hidden_act		,	hidden_dropout_prob=self.hidden_dropout_prob		,	attention_probs_dropout_prob=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	type_vocab_size=self.type_vocab_size		,	initializer_range=self.initializer_range		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Any      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      EsmModel(config=__lowerCamelCase      )
								model.to(__lowerCamelCase      )
								model.eval()
								_lowerCAmelCase    :    str										=      model(__lowerCamelCase		,	attention_mask=__lowerCamelCase      )
								_lowerCAmelCase    :    Optional[Any]										=      model(__lowerCamelCase      )
								_lowerCAmelCase    :    List[Any]										=      model(__lowerCamelCase      )

								self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.hidden_size)      )
								self.parent.assertEqual(result.pooler_output.shape		,	(self.batch_size, self.hidden_size)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		str		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Optional[int]      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      EsmForMaskedLM(config=__lowerCamelCase      )
								model.to(__lowerCamelCase      )
								model.eval()
								_lowerCAmelCase    :    List[Any]										=      model(__lowerCamelCase		,	attention_mask=__lowerCamelCase		,	labels=__lowerCamelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.vocab_size)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		Tuple      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      self.num_labels
								_lowerCAmelCase    :    Optional[int]										=      EsmForTokenClassification(config=__lowerCamelCase      )
								model.to(__lowerCamelCase      )
								model.eval()
								_lowerCAmelCase    :    Optional[Any]										=      model(__lowerCamelCase		,	attention_mask=__lowerCamelCase		,	labels=__lowerCamelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.num_labels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Dict:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      self.prepare_config_and_inputs()
								(
								    _lowerCAmelCase	
								)    :    List[str]										=      config_and_inputs
								_lowerCAmelCase    :    str										=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
								return config, inputs_dict






@require_torch
class        __snake_case  (lowerCamelCase__						,	lowerCamelCase__						,	unittest.TestCase		):
				lowerCAmelCase__				=     False

				lowerCAmelCase__				=     (
				    (
				        EsmForMaskedLM,
				        EsmModel,
				        EsmForSequenceClassification,
				        EsmForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase__				=     ()
				lowerCAmelCase__				=     (
				    {
				        "feature-extraction": EsmModel,
				        "fill-mask": EsmForMaskedLM,
				        "text-classification": EsmForSequenceClassification,
				        "token-classification": EsmForTokenClassification,
				        "zero-shot": EsmForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				lowerCAmelCase__				=     True



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





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      EsmModelTester(self      )
								_lowerCAmelCase    :    int										=      ConfigTester(self		,	config_class=__lowerCamelCase		,	hidden_size=37      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Union[str, Any]      )  ->    Optional[int]:





								'''simple docstring'''




								self.config_tester.run_common_tests()



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__lowerCamelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Any      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
												_lowerCAmelCase    :    List[Any]										=      type
												self.model_tester.create_and_check_model(*__lowerCamelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Tuple:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Tuple:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase      )



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





								'''simple docstring'''




								for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_lowerCAmelCase    :    int										=      EsmModel.from_pretrained(__lowerCamelCase      )
												self.assertIsNotNone(__lowerCamelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      self.model_tester.prepare_config_and_inputs()[0]
								_lowerCAmelCase    :    int										=      EsmEmbeddings(config=__lowerCamelCase      )

								_lowerCAmelCase    :    Optional[int]										=      torch.as_tensor([[12, 31, 13, model.padding_idx]]      )
								_lowerCAmelCase    :    str										=      torch.as_tensor(
								    [
								        [
								            0 + model.padding_idx + 1,
								            1 + model.padding_idx + 1,
								            2 + model.padding_idx + 1,
								            model.padding_idx,
								        ]
								    ]      )
								_lowerCAmelCase    :    List[Any]										=      create_position_ids_from_input_ids(__lowerCamelCase		,	model.padding_idx      )
								self.assertEqual(position_ids.shape		,	expected_positions.shape      )
								self.assertTrue(torch.all(torch.eq(__lowerCamelCase		,	__lowerCamelCase      )      )      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      self.model_tester.prepare_config_and_inputs()[0]
								_lowerCAmelCase    :    Union[str, Any]										=      EsmEmbeddings(config=__lowerCamelCase      )

								_lowerCAmelCase    :    Dict										=      torch.empty(2		,	4		,	30      )
								_lowerCAmelCase    :    Tuple										=      [
								    0 + embeddings.padding_idx + 1,
								    1 + embeddings.padding_idx + 1,
								    2 + embeddings.padding_idx + 1,
								    3 + embeddings.padding_idx + 1,
								]
								_lowerCAmelCase    :    Tuple										=      torch.as_tensor([expected_single_positions, expected_single_positions]      )
								_lowerCAmelCase    :    Optional[int]										=      embeddings.create_position_ids_from_inputs_embeds(__lowerCamelCase      )
								self.assertEqual(position_ids.shape		,	expected_positions.shape      )
								self.assertTrue(torch.all(torch.eq(__lowerCamelCase		,	__lowerCamelCase      )      )      )



				@unittest.skip("""Esm does not support embedding resizing"""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    str:





								'''simple docstring'''




								pass



				@unittest.skip("""Esm does not support embedding resizing"""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    int:





								'''simple docstring'''




								pass



				@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]      )  ->    Optional[Any]:





								'''simple docstring'''




								pass






@require_torch
class        __snake_case  (lowerCamelCase__		):



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Optional[int]:





								'''simple docstring'''




								with torch.no_grad():
												_lowerCAmelCase    :    Tuple										=      EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""      )
												model.eval()
												_lowerCAmelCase    :    Optional[Any]										=      torch.tensor([[0, 1, 2, 3, 4, 5]]      )
												_lowerCAmelCase    :    List[str]										=      model(__lowerCamelCase      )[0]

												_lowerCAmelCase    :    Optional[int]										=      33

												_lowerCAmelCase    :    Tuple										=      torch.Size((1, 6, vocab_size)      )
												self.assertEqual(output.shape		,	__lowerCamelCase      )

												_lowerCAmelCase    :    List[Any]										=      torch.tensor(
												    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]      )
												self.assertTrue(torch.allclose(output[:, :3, :3]		,	__lowerCamelCase		,	atol=1E-4      )      )



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Optional[int]:





								'''simple docstring'''




								with torch.no_grad():
												_lowerCAmelCase    :    Optional[int]										=      EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""      )
												model.eval()

												_lowerCAmelCase    :    Optional[Any]										=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]      )
												_lowerCAmelCase    :    Tuple										=      model(__lowerCamelCase      )[0]
												# compare the actual values for a slice.
												_lowerCAmelCase    :    int										=      torch.tensor(
												    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]      )
												self.assertTrue(torch.allclose(output[:, :3, :3]		,	__lowerCamelCase		,	atol=1E-4      )      )







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        __snake_case  (unittest.TestCase		):



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple		,	_UpperCAmelCase      :		List[str]      )  ->    List[str]:





								'''simple docstring'''




								for model_result in results.values():
												for batch_size, sequence_length in zip(model_result["""bs"""]		,	model_result["""ss"""]      ):
																_lowerCAmelCase    :    List[Any]										=      model_result["""result"""][batch_size][sequence_length]
																self.assertIsNotNone(_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      """sshleifer/tiny-gpt2"""
								_lowerCAmelCase    :    Dict										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    str										=      PyTorchBenchmark(_UpperCAmelCase      )
								_lowerCAmelCase    :    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       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      """sgugger/tiny-distilbert-classification"""
								_lowerCAmelCase    :    Tuple										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	only_pretrain_model=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    List[Any]										=      PyTorchBenchmark(_UpperCAmelCase      )
								_lowerCAmelCase    :    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       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      """sshleifer/tiny-gpt2"""
								_lowerCAmelCase    :    Union[str, Any]										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	torchscript=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    List[str]										=      PyTorchBenchmark(_UpperCAmelCase      )
								_lowerCAmelCase    :    Optional[Any]										=      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       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    str:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      """sshleifer/tiny-gpt2"""
								_lowerCAmelCase    :    Optional[Any]										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	fpaa=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    Union[str, Any]										=      PyTorchBenchmark(_UpperCAmelCase      )
								_lowerCAmelCase    :    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       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    int:





								'''simple docstring'''




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



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    str:





								'''simple docstring'''




								_lowerCAmelCase    :    str										=      """sshleifer/tiny-gpt2"""
								_lowerCAmelCase    :    List[str]										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    List[str]										=      PyTorchBenchmark(_UpperCAmelCase      )
								_lowerCAmelCase    :    Union[str, Any]										=      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       SCREAMING_SNAKE_CASE       (      self      :		Any      )  ->    List[Any]:





								'''simple docstring'''




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



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





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      """sshleifer/tiny-gpt2"""
								_lowerCAmelCase    :    List[Any]										=      AutoConfig.from_pretrained(_UpperCAmelCase      )
								_lowerCAmelCase    :    List[Any]										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    List[str]										=      PyTorchBenchmark(_UpperCAmelCase		,	configs=[config]      )
								_lowerCAmelCase    :    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       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Union[str, Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      """sshleifer/tinier_bart"""
								_lowerCAmelCase    :    Tuple										=      AutoConfig.from_pretrained(_UpperCAmelCase      )
								_lowerCAmelCase    :    Dict										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    Tuple										=      PyTorchBenchmark(_UpperCAmelCase		,	configs=[config]      )
								_lowerCAmelCase    :    Any										=      benchmark.run()
								self.check_results_dict_not_empty(results.time_inference_result      )
								self.check_results_dict_not_empty(results.memory_inference_result      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]      )  ->    List[Any]:





								'''simple docstring'''




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



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





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      """sshleifer/tinier_bart"""
								_lowerCAmelCase    :    Dict										=      AutoConfig.from_pretrained(_UpperCAmelCase      )
								_lowerCAmelCase    :    Any										=      PyTorchBenchmarkArguments(
								    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	multi_process=_UpperCAmelCase		,	)
								_lowerCAmelCase    :    Optional[Any]										=      PyTorchBenchmark(_UpperCAmelCase		,	configs=[config]      )
								_lowerCAmelCase    :    int										=      benchmark.run()
								self.check_results_dict_not_empty(results.time_train_result      )
								self.check_results_dict_not_empty(results.memory_train_result      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    int:





								'''simple docstring'''




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



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





								'''simple docstring'''




								_lowerCAmelCase    :    str										=      """sshleifer/tiny-gpt2"""

								def _check_summary_is_not_empty(_UpperCAmelCase      :		int      ):
												self.assertTrue(hasattr(_UpperCAmelCase		,	"""sequential"""      )      )
												self.assertTrue(hasattr(_UpperCAmelCase		,	"""cumulative"""      )      )
												self.assertTrue(hasattr(_UpperCAmelCase		,	"""current"""      )      )
												self.assertTrue(hasattr(_UpperCAmelCase		,	"""total"""      )      )

								with tempfile.TemporaryDirectory() as tmp_dir:
												_lowerCAmelCase    :    Any										=      PyTorchBenchmarkArguments(
												    models=[MODEL_ID]		,	training=_UpperCAmelCase		,	inference=_UpperCAmelCase		,	sequence_lengths=[8]		,	batch_sizes=[1]		,	log_filename=os.path.join(_UpperCAmelCase		,	"""log.txt"""      )		,	log_print=_UpperCAmelCase		,	trace_memory_line_by_line=_UpperCAmelCase		,	multi_process=_UpperCAmelCase		,	)
												_lowerCAmelCase    :    str										=      PyTorchBenchmark(_UpperCAmelCase      )
												_lowerCAmelCase    :    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(_UpperCAmelCase		,	"""log.txt"""      )      ).exists()      )

