




import argparse
import logging
import os
import re

import tensorflow as tf

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


_A     =			logging.getLogger(__name__)

_A     =			tf.data.AUTOTUNE





def __UpperCamelCase			(       ):
						lowerCAmelCase_      			=  argparse.ArgumentParser(description='''Train a masked language model on TPU.'''  )
						parser.add_argument(
						    '''--pretrained_model_config''' ,  type=_A ,  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=_A ,  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=_A ,  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=_A ,  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=_A ,  help='''Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes.''' ,  )

						parser.add_argument(
						    '''--gcp_project''' ,  type=_A ,  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=_A ,  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=_A ,  default=2**18 ,  help='''Size of the shuffle buffer (in samples)''' ,  )

						parser.add_argument(
						    '''--eval_dataset''' ,  type=_A ,  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=_A ,  default=1 ,  help='''Number of epochs to train for.''' ,  )

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

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

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

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

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

						lowerCAmelCase_      			=  parser.parse_args()
						return args





def __UpperCamelCase			(       _A  ):
						try:
												if args.tpu_name:
																		lowerCAmelCase_      			=  tf.distribute.cluster_resolver.TPUClusterResolver(
																		    args.tpu_name ,  zone=args.tpu_zone ,  project=args.gcp_project  )
												else:
																		lowerCAmelCase_      			=  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(_A  )
						tf.tpu.experimental.initialize_tpu_system(_A  )

						return tpu





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  0
						for file in file_list:
												lowerCAmelCase_      			=  file.split('''/'''  )[-1]
												lowerCAmelCase_      			=  re.search(r'''-\d+-(\d+)\.tfrecord''' ,  _A  ).group(1  )
												lowerCAmelCase_      			=  int(_A  )
												num_samples += sample_count

						return num_samples





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





def __UpperCamelCase			(       _A  ):
						if not args.no_tpu:
												lowerCAmelCase_      			=  initialize_tpu(_A  )
												lowerCAmelCase_      			=  tf.distribute.TPUStrategy(_A  )
						else:
												lowerCAmelCase_      			=  tf.distribute.OneDeviceStrategy(device='''/gpu:0'''  )

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

						lowerCAmelCase_      			=  AutoTokenizer.from_pretrained(args.tokenizer  )
						lowerCAmelCase_      			=  AutoConfig.from_pretrained(args.pretrained_model_config  )
						lowerCAmelCase_      			=  tokenizer.vocab_size

						lowerCAmelCase_      			=  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}."  )
						lowerCAmelCase_      			=  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}."  )

						lowerCAmelCase_      			=  count_samples(_A  )

						lowerCAmelCase_      			=  num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
						lowerCAmelCase_      			=  steps_per_epoch * args.num_epochs

						with strategy.scope():
												lowerCAmelCase_      			=  TFAutoModelForMaskedLM.from_config(_A  )
												model(model.dummy_inputs  )  # Pass some dummy inputs through the model to ensure all the weights are built
												lowerCAmelCase_			,					lowerCAmelCase_      			=  create_optimizer(
												    num_train_steps=_A ,  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=_A ,  metrics=['''accuracy''']  )

						def decode_fn(_A  ):
												lowerCAmelCase_      			=  {
												    '''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(_A ,  _A  )

						# Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
						# use their methods in our data pipeline.
						lowerCAmelCase_      			=  DataCollatorForLanguageModeling(
						    tokenizer=_A ,  mlm_probability=args.mlm_probability ,  mlm=_A ,  return_tensors='''tf'''  )

						def mask_with_collator(_A  ):
												# TF really needs an isin() function
												lowerCAmelCase_      			=  (
												    ~tf.cast(batch['''attention_mask'''] ,  tf.bool  )
												    | (batch['''input_ids'''] == tokenizer.cls_token_id)
												    | (batch['''input_ids'''] == tokenizer.sep_token_id)
												)
												lowerCAmelCase_			,					lowerCAmelCase_      			=  data_collator.tf_mask_tokens(
												    batch['''input_ids'''] ,  vocab_size=len(_A  ) ,  mask_token_id=tokenizer.mask_token_id ,  special_tokens_mask=_A ,  )
												return batch

						lowerCAmelCase_      			=  args.per_replica_batch_size * strategy.num_replicas_in_sync

						lowerCAmelCase_      			=  prepare_dataset(
						    _A ,  decode_fn=_A ,  mask_fn=_A ,  batch_size=_A ,  shuffle=_A ,  shuffle_buffer_size=args.shuffle_buffer_size ,  )

						lowerCAmelCase_      			=  prepare_dataset(
						    _A ,  decode_fn=_A ,  mask_fn=_A ,  batch_size=_A ,  shuffle=_A ,  )

						lowerCAmelCase_      			=  []
						if args.hub_model_id:
												callbacks.append(
												    PushToHubCallback(output_dir=args.output_dir ,  hub_model_id=args.hub_model_id ,  tokenizer=_A  )  )

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

						model.save_pretrained(args.output_dir  )


if __name__ == "__main__":
						_A     =			parse_args()
						main(args)










from functools import lru_cache





@lru_cache
def __UpperCamelCase			(       _A  ):
						if num < 0:
												raise ValueError('''Number should not be negative.'''  )

						return 1 if num in (0, 1) else num * factorial(num - 1  )


if __name__ == "__main__":
						import doctest

						doctest.testmod()




