



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		):







						"""simple docstring"""




						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase  ) ->				Dict:
									"""simple docstring"""




									for model_result in results.values():
												for batch_size, sequence_length in zip(model_result["""bs"""] ,       model_result["""ss"""]  ):
															a__       :						Optional[Any]							=  model_result["""result"""][batch_size][sequence_length]
															self.assertIsNotNone(__lowercase  )

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
									"""simple docstring"""




									a__       :						Tuple							=  """sshleifer/tiny-gpt2"""
									a__       :						str							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						List[str]							=  PyTorchBenchmark(__lowercase  )
									a__       :						Dict							=  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  ) ->				List[str]:
									"""simple docstring"""




									a__       :						Any							=  """sgugger/tiny-distilbert-classification"""
									a__       :						Any							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       only_pretrain_model=__lowercase ,       )
									a__       :						str							=  PyTorchBenchmark(__lowercase  )
									a__       :						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]:
									"""simple docstring"""




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

						@unittest.skipIf(torch_device == """cpu""" ,       """Cant do half precision"""  )
						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
									"""simple docstring"""




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




									a__       :						List[str]							=  """sshleifer/tiny-gpt2"""
									a__       :						int							=  AutoConfig.from_pretrained(__lowercase  )
									# set architectures equal to `None`
									a__       :						List[Any]							=  None
									a__       :						int							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						int							=  PyTorchBenchmark(__lowercase ,       configs=[config]  )
									a__       :						Tuple							=  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  ) ->				List[str]:
									"""simple docstring"""




									a__       :						Dict							=  """sshleifer/tiny-gpt2"""
									a__       :						int							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						Optional[Any]							=  PyTorchBenchmark(__lowercase  )
									a__       :						Optional[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:
									"""simple docstring"""




									a__       :						Any							=  """sshleifer/tiny-gpt2"""
									a__       :						Optional[Any]							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       fpaa=__lowercase ,       multi_process=__lowercase ,       )
									a__       :						Any							=  PyTorchBenchmark(__lowercase  )
									a__       :						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  ) ->				Any:
									"""simple docstring"""




									a__       :						List[Any]							=  """sshleifer/tiny-gpt2"""
									a__       :						str							=  AutoConfig.from_pretrained(__lowercase  )
									a__       :						Dict							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						Any							=  PyTorchBenchmark(__lowercase ,       configs=[config]  )
									a__       :						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  ) ->				Tuple:
									"""simple docstring"""




									a__       :						Optional[Any]							=  """sshleifer/tinier_bart"""
									a__       :						List[str]							=  AutoConfig.from_pretrained(__lowercase  )
									a__       :						str							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						List[str]							=  PyTorchBenchmark(__lowercase ,       configs=[config]  )
									a__       :						Dict							=  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  ) ->				Tuple:
									"""simple docstring"""




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




									a__       :						Tuple							=  """sshleifer/tinier_bart"""
									a__       :						Optional[Any]							=  AutoConfig.from_pretrained(__lowercase  )
									a__       :						str							=  PyTorchBenchmarkArguments(
									    models=[MODEL_ID] ,       training=__lowercase ,       inference=__lowercase ,       sequence_lengths=[8] ,       batch_sizes=[1] ,       multi_process=__lowercase ,       )
									a__       :						Union[str, Any]							=  PyTorchBenchmark(__lowercase ,       configs=[config]  )
									a__       :						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  ) ->				str:
									"""simple docstring"""




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

						def 							SCREAMING_SNAKE_CASE__(       self  ) ->				str:
									"""simple docstring"""




									a__       :						List[str]							=  """sshleifer/tiny-gpt2"""

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

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









import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_lowercase		:				Union[str, Any]      		=data_utils.TransfoXLTokenizer
_lowercase		:				Tuple      		=data_utils.TransfoXLCorpus
_lowercase		:				Optional[int]      		=data_utils
_lowercase		:				int      		=data_utils

def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       List[Any]						,       _lowercase    :       Optional[Any]						,       _lowercase    :       Optional[Any])							->	List[Any]:







			"""simple docstring"""
			if transfo_xl_dataset_file:
						# Convert a pre-processed corpus (see original TensorFlow repo)
						with open(_lowercase						,       """rb""") as fp:
									a__       :						Optional[Any]							=  pickle.load(_lowercase						,       encoding="""latin1""")
						# Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
						a__       :						Dict							=  pytorch_dump_folder_path + """/""" + VOCAB_FILES_NAMES["""pretrained_vocab_file"""]
						print(F'''Save vocabulary to {pytorch_vocab_dump_path}''')
						a__       :						List[str]							=  corpus.vocab.__dict__
						torch.save(_lowercase						,       _lowercase)

						a__       :						Optional[int]							=  corpus.__dict__
						corpus_dict_no_vocab.pop("""vocab"""						,       _lowercase)
						a__       :						Tuple							=  pytorch_dump_folder_path + """/""" + CORPUS_NAME
						print(F'''Save dataset to {pytorch_dataset_dump_path}''')
						torch.save(_lowercase						,       _lowercase)

			if tf_checkpoint_path:
						# Convert a pre-trained TensorFlow model
						a__       :						Optional[Any]							=  os.path.abspath(_lowercase)
						a__       :						str							=  os.path.abspath(_lowercase)

						print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.''')
						# Initialise PyTorch model
						if transfo_xl_config_file == "":
									a__       :						List[str]							=  TransfoXLConfig()
						else:
									a__       :						Any							=  TransfoXLConfig.from_json_file(_lowercase)
						print(F'''Building PyTorch model from configuration: {config}''')
						a__       :						Any							=  TransfoXLLMHeadModel(_lowercase)

						a__       :						Optional[Any]							=  load_tf_weights_in_transfo_xl(_lowercase						,       _lowercase						,       _lowercase)
						# Save pytorch-model
						a__       :						Tuple							=  os.path.join(_lowercase						,       _lowercase)
						a__       :						Optional[int]							=  os.path.join(_lowercase						,       _lowercase)
						print(F'''Save PyTorch model to {os.path.abspath(_lowercase)}''')
						torch.save(model.state_dict()						,       _lowercase)
						print(F'''Save configuration file to {os.path.abspath(_lowercase)}''')
						with open(_lowercase						,       """w"""						,       encoding="""utf-8""") as f:
									f.write(config.to_json_string())


if __name__ == "__main__":
		_lowercase		:				Tuple      		=argparse.ArgumentParser()
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default=None,
		    type=str,
		    required=True,
		    help="Path to the folder to store the PyTorch model or dataset/vocab.",
		)
		parser.add_argument(
		    "--tf_checkpoint_path",
		    default="",
		    type=str,
		    help="An optional path to a TensorFlow checkpoint path to be converted.",
		)
		parser.add_argument(
		    "--transfo_xl_config_file",
		    default="",
		    type=str,
		    help=(
		        "An optional config json file corresponding to the pre-trained BERT model. \n"
		        "This specifies the model architecture."
		    ),
		)
		parser.add_argument(
		    "--transfo_xl_dataset_file",
		    default="",
		    type=str,
		    help="An optional dataset file to be converted in a vocabulary.",
		)
		_lowercase		:				Dict      		=parser.parse_args()
		convert_transfo_xl_checkpoint_to_pytorch(
		    args.tf_checkpoint_path,
		    args.transfo_xl_config_file,
		    args.pytorch_dump_folder_path,
		    args.transfo_xl_dataset_file,
		)




