




import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


UpperCAmelCase_		:				List[Any]									=     logging.getLogger(__name__)

UpperCAmelCase_		:				int									=     tf.data.AUTOTUNE




def 		lowerCAmelCase_   (	):
				__magic_name__	:	Optional[Any]								=argparse.ArgumentParser(description="""Train a masked language model on TPU."""   )
				parser.add_argument(
				    """--pretrained_model_config"""	,     type=lowerCamelCase	,     default="""roberta-base"""	,     help="""The model config to use. Note that we don't copy the model's weights, only the config!"""	,     )
				parser.add_argument(
				    """--tokenizer"""	,     type=lowerCamelCase	,     default="""unigram-tokenizer-wikitext"""	,     help="""The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model's vocab size."""	,     )

				parser.add_argument(
				    """--per_replica_batch_size"""	,     type=lowerCamelCase	,     default=8	,     help="""Batch size per TPU core."""	,     )

				parser.add_argument(
				    """--no_tpu"""	,     action="""store_true"""	,     help="""If set, run on CPU and don't try to initialize a TPU. Useful for debugging on non-TPU instances."""	,     )

				parser.add_argument(
				    """--tpu_name"""	,     type=lowerCamelCase	,     help="""Name of TPU resource to initialize. Should be blank on Colab, and 'local' on TPU VMs."""	,     default="""local"""	,     )

				parser.add_argument(
				    """--tpu_zone"""	,     type=lowerCamelCase	,     help="""Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."""	,     )

				parser.add_argument(
				    """--gcp_project"""	,     type=lowerCamelCase	,     help="""Google cloud project name. Only used for non-Colab TPU nodes."""   )

				parser.add_argument(
				    """--bfloat16"""	,     action="""store_true"""	,     help="""Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU."""	,     )

				parser.add_argument(
				    """--train_dataset"""	,     type=lowerCamelCase	,     help="""Path to training dataset to load. If the path begins with `gs://`"""
				    """ then the dataset will be loaded from a Google Cloud Storage bucket."""	,     )

				parser.add_argument(
				    """--shuffle_buffer_size"""	,     type=lowerCamelCase	,     default=2**18	,     help="""Size of the shuffle buffer (in samples)"""	,     )

				parser.add_argument(
				    """--eval_dataset"""	,     type=lowerCamelCase	,     help="""Path to evaluation dataset to load. If the path begins with `gs://`"""
				    """ then the dataset will be loaded from a Google Cloud Storage bucket."""	,     )

				parser.add_argument(
				    """--num_epochs"""	,     type=lowerCamelCase	,     default=1	,     help="""Number of epochs to train for."""	,     )

				parser.add_argument(
				    """--learning_rate"""	,     type=lowerCamelCase	,     default=1E-4	,     help="""Learning rate to use for training."""	,     )

				parser.add_argument(
				    """--weight_decay_rate"""	,     type=lowerCamelCase	,     default=1E-3	,     help="""Weight decay rate to use for training."""	,     )

				parser.add_argument(
				    """--max_length"""	,     type=lowerCamelCase	,     default=512	,     help="""Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py"""	,     )

				parser.add_argument(
				    """--mlm_probability"""	,     type=lowerCamelCase	,     default=0.1_5	,     help="""Fraction of tokens to mask during training."""	,     )

				parser.add_argument("""--output_dir"""	,     type=lowerCamelCase	,     required=lowerCamelCase	,     help="""Path to save model checkpoints to."""   )
				parser.add_argument("""--hub_model_id"""	,     type=lowerCamelCase	,     help="""Model ID to upload to on the Hugging Face Hub."""   )

				__magic_name__	:	int								=parser.parse_args()
				return args




def 		lowerCAmelCase_   (	lowerCamelCase   ):
				try:
								if args.tpu_name:
												__magic_name__	:	Tuple								=tf.distribute.cluster_resolver.TPUClusterResolver(
												    args.tpu_name	,     zone=args.tpu_zone	,     project=args.gcp_project   )
								else:
												__magic_name__	:	List[Any]								=tf.distribute.cluster_resolver.TPUClusterResolver()
				except ValueError:
								raise RuntimeError(
								    """Couldn't connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or """
								    """--gcp_project. When running on a TPU VM, use --tpu_name local."""   )

				tf.config.experimental_connect_to_cluster(lowerCamelCase   )
				tf.tpu.experimental.initialize_tpu_system(lowerCamelCase   )

				return tpu




def 		lowerCAmelCase_   (	lowerCamelCase   ):
				__magic_name__	:	Tuple								=0
				for file in file_list:
								__magic_name__	:	Any								=file.split("""/"""   )[-1]
								__magic_name__	:	Union[str, Any]								=re.search(R"""-\d+-(\d+)\.tfrecord"""	,     lowerCamelCase   ).group(1   )
								__magic_name__	:	Any								=int(lowerCamelCase   )
								num_samples += sample_count

				return num_samples




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase=None   ):
				__magic_name__	:	Optional[Any]								=count_samples(lowerCamelCase   )
				__magic_name__	:	str								=tf.data.Dataset.from_tensor_slices(lowerCamelCase   )
				if shuffle:
								__magic_name__	:	Optional[Any]								=dataset.shuffle(len(lowerCamelCase   )   )
				__magic_name__	:	str								=tf.data.TFRecordDataset(lowerCamelCase	,     num_parallel_reads=lowerCamelCase   )
				# TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
				__magic_name__	:	List[str]								=dataset.apply(tf.data.experimental.assert_cardinality(lowerCamelCase   )   )
				__magic_name__	:	Union[str, Any]								=dataset.map(lowerCamelCase	,     num_parallel_calls=lowerCamelCase   )
				if shuffle:
								assert shuffle_buffer_size is not None
								__magic_name__	:	int								=dataset.shuffle(args.shuffle_buffer_size   )
				__magic_name__	:	int								=dataset.batch(lowerCamelCase	,     drop_remainder=lowerCamelCase   )
				__magic_name__	:	List[Any]								=dataset.map(lowerCamelCase	,     num_parallel_calls=lowerCamelCase   )
				__magic_name__	:	Tuple								=dataset.prefetch(lowerCamelCase   )
				return dataset




def 		lowerCAmelCase_   (	lowerCamelCase   ):
				if not args.no_tpu:
								__magic_name__	:	Any								=initialize_tpu(lowerCamelCase   )
								__magic_name__	:	int								=tf.distribute.TPUStrategy(lowerCamelCase   )
				else:
								__magic_name__	:	Any								=tf.distribute.OneDeviceStrategy(device="""/gpu:0"""   )

				if args.bfloataa:
								tf.keras.mixed_precision.set_global_policy("""mixed_bfloat16"""   )

				__magic_name__	:	Optional[int]								=AutoTokenizer.from_pretrained(args.tokenizer   )
				__magic_name__	:	Optional[int]								=AutoConfig.from_pretrained(args.pretrained_model_config   )
				__magic_name__	:	Tuple								=tokenizer.vocab_size

				__magic_name__	:	Optional[Any]								=tf.io.gfile.glob(os.path.join(args.train_dataset	,     """*.tfrecord"""   )   )
				if not training_records:
								raise ValueError(F"No .tfrecord files found in {args.train_dataset}."   )
				__magic_name__	:	int								=tf.io.gfile.glob(os.path.join(args.eval_dataset	,     """*.tfrecord"""   )   )
				if not eval_records:
								raise ValueError(F"No .tfrecord files found in {args.eval_dataset}."   )

				__magic_name__	:	Optional[Any]								=count_samples(lowerCamelCase   )

				__magic_name__	:	List[str]								=num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
				__magic_name__	:	Any								=steps_per_epoch * args.num_epochs

				with strategy.scope():
								__magic_name__	:	str								=TFAutoModelForMaskedLM.from_config(lowerCamelCase   )
								model(model.dummy_inputs   )  # Pass some dummy inputs through the model to ensure all the weights are built
								__magic_name__	:	Dict								=create_optimizer(
								    num_train_steps=lowerCamelCase	,     num_warmup_steps=total_train_steps // 20	,     init_lr=args.learning_rate	,     weight_decay_rate=args.weight_decay_rate	,     )

								# Transformers models compute the right loss for their task by default when labels are passed, and will
								# use this for training unless you specify your own loss function in compile().
								model.compile(optimizer=lowerCamelCase	,     metrics=["""accuracy"""]   )

				def decode_fn(lowerCamelCase   ):
								__magic_name__	:	Optional[Any]								={
								    """input_ids""": tf.io.FixedLenFeature(dtype=tf.intaa	,     shape=(args.max_length,)   ),
								    """attention_mask""": tf.io.FixedLenFeature(dtype=tf.intaa	,     shape=(args.max_length,)   ),
								}
								return tf.io.parse_single_example(lowerCamelCase	,     lowerCamelCase   )

				# Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
				# use their methods in our data pipeline.
				__magic_name__	:	Union[str, Any]								=DataCollatorForLanguageModeling(
				    tokenizer=lowerCamelCase	,     mlm_probability=args.mlm_probability	,     mlm=lowerCamelCase	,     return_tensors="""tf"""   )

				def mask_with_collator(lowerCamelCase   ):
								# TF really needs an isin() function
								__magic_name__	:	Optional[Any]								=(
								    ~tf.cast(batch["""attention_mask"""]	,     tf.bool   )
								    | (batch["""input_ids"""] == tokenizer.cls_token_id)
								    | (batch["""input_ids"""] == tokenizer.sep_token_id)
								)
								__magic_name__	:	int								=data_collator.tf_mask_tokens(
								    batch["""input_ids"""]	,     vocab_size=len(lowerCamelCase   )	,     mask_token_id=tokenizer.mask_token_id	,     special_tokens_mask=lowerCamelCase	,     )
								return batch

				__magic_name__	:	Dict								=args.per_replica_batch_size * strategy.num_replicas_in_sync

				__magic_name__	:	List[Any]								=prepare_dataset(
				    lowerCamelCase	,     decode_fn=lowerCamelCase	,     mask_fn=lowerCamelCase	,     batch_size=lowerCamelCase	,     shuffle=lowerCamelCase	,     shuffle_buffer_size=args.shuffle_buffer_size	,     )

				__magic_name__	:	Dict								=prepare_dataset(
				    lowerCamelCase	,     decode_fn=lowerCamelCase	,     mask_fn=lowerCamelCase	,     batch_size=lowerCamelCase	,     shuffle=lowerCamelCase	,     )

				__magic_name__	:	int								=[]
				if args.hub_model_id:
								callbacks.append(
								    PushToHubCallback(output_dir=args.output_dir	,     hub_model_id=args.hub_model_id	,     tokenizer=lowerCamelCase   )   )

				model.fit(
				    lowerCamelCase	,     validation_data=lowerCamelCase	,     epochs=args.num_epochs	,     callbacks=lowerCamelCase	,     )

				model.save_pretrained(args.output_dir   )


if __name__ == "__main__":
							UpperCAmelCase_		:				List[str]									=     parse_args()
							main(args)






import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.framework.errors_impl import ResourceExhaustedError

							from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
							import pyanvml.pyanvml as nvml

UpperCAmelCase_		:				Optional[Any]									=     logging.get_logger(__name__)




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase   ):
				def run_func(lowerCamelCase   ):
								@wraps(lowerCamelCase   )
								def run_in_eager_mode(*lowerCamelCase	,     **lowerCamelCase   ):
												return func(*lowerCamelCase	,     **lowerCamelCase   )

								@wraps(lowerCamelCase   )
								@tf.function(experimental_compile=lowerCamelCase   )
								def run_in_graph_mode(*lowerCamelCase	,     **lowerCamelCase   ):
												return func(*lowerCamelCase	,     **lowerCamelCase   )

								if do_eager_mode is True:
												if use_xla is not False:
																raise ValueError(
																    """Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."""   )
												return run_in_eager_mode
								else:
												return run_in_graph_mode

				return run_func




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase   ):
				__magic_name__	:	Dict								=random.Random()
				__magic_name__	:	Union[str, Any]								=[rng.randint(0	,     vocab_size - 1   ) for i in range(batch_size * sequence_length   )]
				return tf.constant(lowerCamelCase	,     shape=(batch_size, sequence_length)	,     dtype=tf.intaa   )




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					42
						UpperCamelCase     						=					42
						UpperCamelCase     						=					"TensorFlow"

						@property
						def     A__						(		self   :str ):
										'''simple docstring'''

										return tf.__version__

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Dict								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Union[str, Any]								=self._prepare_inference_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_speed(_inference )

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Any								=self._prepare_train_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_speed(_train )

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    , __snake_case )
										__magic_name__	:	int								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Tuple								=self._prepare_inference_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_memory(_inference )

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    , __snake_case )
										__magic_name__	:	Tuple								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )

										__magic_name__	:	Any								=self._prepare_train_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_memory(_train )

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	int								=self.config_dict[model_name]

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported.""" )

										__magic_name__	:	Any								=(
										    hasattr(__snake_case    , """architectures""" )
										    and isinstance(config.architectures    , __snake_case )
										    and len(config.architectures ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		__magic_name__	:	Optional[int]								="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		__magic_name__	:	Optional[int]								=__import__("""transformers"""    , fromlist=[model_class] )
																		__magic_name__	:	Optional[Any]								=getattr(__snake_case    , __snake_case )
																		__magic_name__	:	Optional[Any]								=model_cls(__snake_case )
														except ImportError:
																		raise ImportError(
																		    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`.""" )
										else:
														__magic_name__	:	Optional[int]								=TF_MODEL_MAPPING[config.__class__](__snake_case )

										# encoder-decoder has vocab size saved differently
										__magic_name__	:	List[str]								=config.vocab_size if hasattr(__snake_case    , """vocab_size""" ) else config.encoder.vocab_size
										__magic_name__	:	Any								=random_input_ids(__snake_case    , __snake_case    , __snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_decoder_forward():
														return model(__snake_case    , decoder_input_ids=__snake_case    , training=__snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_forward():
														return model(__snake_case    , training=__snake_case )

										__magic_name__	:	Tuple								=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

										return _inference

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=self.config_dict[model_name]

										if self.args.eager_mode is not False:
														raise ValueError("""Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.""" )

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported.""" )

										__magic_name__	:	int								=(
										    hasattr(__snake_case    , """architectures""" )
										    and isinstance(config.architectures    , __snake_case )
										    and len(config.architectures ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		__magic_name__	:	List[Any]								="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		__magic_name__	:	Optional[int]								=__import__("""transformers"""    , fromlist=[model_class] )
																		__magic_name__	:	str								=getattr(__snake_case    , __snake_case )
																		__magic_name__	:	List[Any]								=model_cls(__snake_case )
														except ImportError:
																		raise ImportError(
																		    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`.""" )
										else:
														__magic_name__	:	int								=TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__snake_case )

										# encoder-decoder has vocab size saved differently
										__magic_name__	:	int								=config.vocab_size if hasattr(__snake_case    , """vocab_size""" ) else config.encoder.vocab_size
										__magic_name__	:	Optional[Any]								=random_input_ids(__snake_case    , __snake_case    , __snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_decoder_train():
														__magic_name__	:	List[str]								=model(__snake_case    , decoder_input_ids=__snake_case    , labels=__snake_case    , training=__snake_case )[0]
														__magic_name__	:	int								=tf.gradients(__snake_case    , model.trainable_variables )
														return gradients

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_train():
														__magic_name__	:	str								=model(__snake_case    , labels=__snake_case    , training=__snake_case )[0]
														__magic_name__	:	int								=tf.gradients(__snake_case    , model.trainable_variables )
														return gradients

										__magic_name__	:	Union[str, Any]								=encoder_decoder_train if config.is_encoder_decoder else encoder_train

										return _train

						def     A__						(		self   :Any    , __snake_case   :Union[str, Any] ):
										'''simple docstring'''

										with self.args.strategy.scope():
														try:
																		if self.args.is_tpu or self.args.use_xla:
																						# run additional 10 times to stabilize compilation for tpu
																						logger.info("""Do inference on TPU. Running model 5 times to stabilize compilation""" )
																						timeit.repeat(__snake_case    , repeat=1    , number=5 )

																		# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																		__magic_name__	:	Union[str, Any]								=timeit.repeat(
																		    __snake_case    , repeat=self.args.repeat    , number=10    , )

																		return min(__snake_case ) / 10.0
														except ResourceExhaustedError as e:
																		self.print_fn(f"Doesn't fit on GPU. {e}" )


						def     A__						(		self   :Any    , __snake_case   :Callable[[], None] ):
										'''simple docstring'''

										logger.info(
										    """Note that TensorFlow allocates more memory than """
										    """it might need to speed up computation. """
										    """The memory reported here corresponds to the memory """
										    """reported by `nvidia-smi`, which can vary depending """
										    """on total available memory on the GPU that is used.""" )
										with self.args.strategy.scope():
														try:
																		if self.args.trace_memory_line_by_line:
																						if not self.args.eager_mode:
																										raise ValueError(
																										    """`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"""
																										    """ consumption line by line.""" )
																						__magic_name__	:	Union[str, Any]								=start_memory_tracing("""transformers""" )

																		if self.args.is_tpu:
																						# tpu
																						raise NotImplementedError(
																						    """Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking"""
																						    """ with `args.memory=False`""" )
																		elif self.args.is_gpu:
																						# gpu
																						if not is_pyanvml_available():
																										logger.warning(
																										    """py3nvml not installed, we won't log GPU memory usage. """
																										    """Install py3nvml (pip install py3nvml) to log information about GPU.""" )
																										__magic_name__	:	str								="""N/A"""
																						else:
																										logger.info(
																										    """Measuring total GPU usage on GPU device. Make sure to not have additional processes"""
																										    """ running on the same GPU.""" )
																										# init nvml
																										nvml.nvmlInit()
																										func()
																										__magic_name__	:	List[str]								=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
																										__magic_name__	:	Tuple								=nvml.nvmlDeviceGetMemoryInfo(__snake_case )
																										__magic_name__	:	Any								=meminfo.used
																										__magic_name__	:	str								=Memory(__snake_case )
																										# shutdown nvml
																										nvml.nvmlShutdown()
																		else:
																						# cpu
																						if self.args.trace_memory_line_by_line:
																										logger.info(
																										    """When enabling line by line tracing, the max peak memory for CPU is inaccurate in"""
																										    """ TensorFlow.""" )
																										__magic_name__	:	List[str]								=None
																						else:
																										__magic_name__	:	List[Any]								=measure_peak_memory_cpu(__snake_case )
																										__magic_name__	:	str								=Memory(__snake_case ) if isinstance(__snake_case    , __snake_case ) else memory_bytes
																		if self.args.trace_memory_line_by_line:
																						__magic_name__	:	List[Any]								=stop_memory_tracing(__snake_case )
																						if memory is None:
																										__magic_name__	:	Any								=summary.total
																		else:
																						__magic_name__	:	Optional[Any]								=None

																		return memory, summary
														except ResourceExhaustedError as e:
																		self.print_fn(f"Doesn't fit on GPU. {e}" )
																		return "N/A", None




