





import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
	import tensorflow as tf

	from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments




@require_tf
class 	UpperCamelCase_							(  unittest.TestCase     ):


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



									"""simple docstring"""



									for model_result in results.values():
											for batch_size, sequence_length in zip(model_result["""bs"""]				,					model_result["""ss"""]		):
													SCREAMING_SNAKE_CASE__								=   model_result["""result"""][batch_size][sequence_length]
													self.assertIsNotNone(__A		)


							def   _snake_case	(				self     :Any		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					eager_mode=__A				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sgugger/tiny-distilbert-classification"""
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					only_pretrain_model=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


							def   _snake_case	(				self     :Tuple		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   AutoConfig.from_pretrained(__A		)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					eager_mode=__A				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A				,					[config]		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   AutoConfig.from_pretrained(__A		)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A				,					[config]		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


							def   _snake_case	(				self     :Dict		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_train_result		)
									self.check_results_dict_not_empty(results.memory_train_result		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   AutoConfig.from_pretrained(__A		)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A				,					[config]		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_train_result		)
									self.check_results_dict_not_empty(results.memory_train_result		)


							def   _snake_case	(				self     :int		) -> Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """patrickvonplaten/t5-tiny-random"""
									SCREAMING_SNAKE_CASE__								=   AutoConfig.from_pretrained(__A		)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A				,					configs=[config]		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


							@unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices("""GPU"""		)		) == 0				,					"""Cannot do xla on CPU."""		)
							def   _snake_case	(				self     :List[str]		) -> List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
									    models=[MODEL_ID]				,					training=__A				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					use_xla=__A				,					multi_process=__A				,					)
									SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
									SCREAMING_SNAKE_CASE__								=   benchmark.run()
									self.check_results_dict_not_empty(results.time_inference_result		)
									self.check_results_dict_not_empty(results.memory_inference_result		)


							def   _snake_case	(				self     :Dict		) -> Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""
									with tempfile.TemporaryDirectory() as tmp_dir:
											SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
											    models=[MODEL_ID]				,					inference=__A				,					save_to_csv=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					inference_time_csv_file=os.path.join(__A				,					"""inf_time.csv"""		)				,					inference_memory_csv_file=os.path.join(__A				,					"""inf_mem.csv"""		)				,					env_info_csv_file=os.path.join(__A				,					"""env.csv"""		)				,					multi_process=__A				,					)
											SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
											benchmark.run()
											self.assertTrue(Path(os.path.join(__A				,					"""inf_time.csv"""		)		).exists()		)
											self.assertTrue(Path(os.path.join(__A				,					"""inf_mem.csv"""		)		).exists()		)
											self.assertTrue(Path(os.path.join(__A				,					"""env.csv"""		)		).exists()		)


							def   _snake_case	(				self     :Dict		) -> Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """sshleifer/tiny-gpt2"""

									def _check_summary_is_not_empty(__A     :Optional[Any]		):
											self.assertTrue(hasattr(__A				,					"""sequential"""		)		)
											self.assertTrue(hasattr(__A				,					"""cumulative"""		)		)
											self.assertTrue(hasattr(__A				,					"""current"""		)		)
											self.assertTrue(hasattr(__A				,					"""total"""		)		)

									with tempfile.TemporaryDirectory() as tmp_dir:
											SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmarkArguments(
											    models=[MODEL_ID]				,					inference=__A				,					sequence_lengths=[8]				,					batch_sizes=[1]				,					log_filename=os.path.join(__A				,					"""log.txt"""		)				,					log_print=__A				,					trace_memory_line_by_line=__A				,					eager_mode=__A				,					multi_process=__A				,					)
											SCREAMING_SNAKE_CASE__								=   TensorFlowBenchmark(__A		)
											SCREAMING_SNAKE_CASE__								=   benchmark.run()
											_check_summary_is_not_empty(result.inference_summary		)
											self.assertTrue(Path(os.path.join(__A				,					"""log.txt"""		)		).exists()		)






import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
_lowerCamelCase 					=      logging.getLogger(__name__)

if __name__ == "__main__":
	_lowerCamelCase 					=      argparse.ArgumentParser(
	    description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
	)
	parser.add_argument(
	    '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
	)
	parser.add_argument(
	    '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
	)
	parser.add_argument('--vocab_size', default=30522, type=int)
	_lowerCamelCase 					=      parser.parse_args()

	logger.info(F'''Loading data from {args.data_file}''')
	with open(args.data_file, 'rb') as fp:
		_lowerCamelCase 					=      pickle.load(fp)

	logger.info('Counting occurrences for MLM.')
	_lowerCamelCase 					=      Counter()
	for tk_ids in data:
		counter.update(tk_ids)
	_lowerCamelCase 					=      [0] * args.vocab_size
	for k, v in counter.items():
		_lowerCamelCase 					=      v

	logger.info(F'''Dump to {args.token_counts_dump}''')
	with open(args.token_counts_dump, 'wb') as handle:
		pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)