




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

     '''simple docstring'''


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

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

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

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

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

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

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

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

     UpperCAmelCase_     :				Optional[int]  =     parser.parse_args()
     return args







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     try:
          if args.tpu_name:
               UpperCAmelCase_     :				List[Any]  =     tf.distribute.cluster_resolver.TPUClusterResolver(
                   args.tpu_name					,		zone=args.tpu_zone					,		project=args.gcp_project )
          else:
               UpperCAmelCase_     :				Optional[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(_lowercase )
     tf.tpu.experimental.initialize_tpu_system(_lowercase )

     return tpu







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				List[str]  =     0
     for file in file_list:
          UpperCAmelCase_     :				Optional[int]  =     file.split('''/''' )[-1]
          UpperCAmelCase_     :				Tuple  =     re.search(r'''-\d+-(\d+)\.tfrecord'''					,		_lowercase ).group(1 )
          UpperCAmelCase_     :				Dict  =     int(_lowercase )
          num_samples += sample_count

     return num_samples







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase=None ):

     '''simple docstring'''


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







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     if not args.no_tpu:
          UpperCAmelCase_     :				List[Any]  =     initialize_tpu(_lowercase )
          UpperCAmelCase_     :				Optional[int]  =     tf.distribute.TPUStrategy(_lowercase )
     else:
          UpperCAmelCase_     :				Dict  =     tf.distribute.OneDeviceStrategy(device='''/gpu:0''' )

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

     UpperCAmelCase_     :				str  =     AutoTokenizer.from_pretrained(args.tokenizer )
     UpperCAmelCase_     :				List[str]  =     AutoConfig.from_pretrained(args.pretrained_model_config )
     UpperCAmelCase_     :				int  =     tokenizer.vocab_size

     UpperCAmelCase_     :				Union[str, 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}.''' )
     UpperCAmelCase_     :				Tuple  =     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}.''' )

     UpperCAmelCase_     :				Dict  =     count_samples(_lowercase )

     UpperCAmelCase_     :				Any  =     num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
     UpperCAmelCase_     :				Optional[int]  =     steps_per_epoch * args.num_epochs

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

     def decode_fn(_lowercase ):
          UpperCAmelCase_     :				List[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(_lowercase					,		_lowercase )

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

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

     UpperCAmelCase_     :				Any  =     args.per_replica_batch_size * strategy.num_replicas_in_sync

     UpperCAmelCase_     :				Any  =     prepare_dataset(
         _lowercase					,		decode_fn=_lowercase					,		mask_fn=_lowercase					,		batch_size=_lowercase					,		shuffle=_lowercase					,		shuffle_buffer_size=args.shuffle_buffer_size					,		)

     UpperCAmelCase_     :				Dict  =     prepare_dataset(
         _lowercase					,		decode_fn=_lowercase					,		mask_fn=_lowercase					,		batch_size=_lowercase					,		shuffle=_lowercase					,		)

     UpperCAmelCase_     :				Optional[Any]  =     []
     if args.hub_model_id:
          callbacks.append(
              PushToHubCallback(output_dir=args.output_dir					,		hub_model_id=args.hub_model_id					,		tokenizer=_lowercase ) )

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

     model.save_pretrained(args.output_dir )


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





from functools import reduce

__a							     =					(
    '73167176531330624919225119674426574742355349194934'
    '96983520312774506326239578318016984801869478851843'
    '85861560789112949495459501737958331952853208805511'
    '12540698747158523863050715693290963295227443043557'
    '66896648950445244523161731856403098711121722383113'
    '62229893423380308135336276614282806444486645238749'
    '30358907296290491560440772390713810515859307960866'
    '70172427121883998797908792274921901699720888093776'
    '65727333001053367881220235421809751254540594752243'
    '52584907711670556013604839586446706324415722155397'
    '53697817977846174064955149290862569321978468622482'
    '83972241375657056057490261407972968652414535100474'
    '82166370484403199890008895243450658541227588666881'
    '16427171479924442928230863465674813919123162824586'
    '17866458359124566529476545682848912883142607690042'
    '24219022671055626321111109370544217506941658960408'
    '07198403850962455444362981230987879927244284909188'
    '84580156166097919133875499200524063689912560717606'
    '05886116467109405077541002256983155200055935729725'
    '71636269561882670428252483600823257530420752963450'
)







def lowerCamelCase__  (     _lowercase = N ):

     '''simple docstring'''


     return max(
         # mypy cannot properly interpret reduce
         int(reduce(lambda _lowercase					,		_lowercase : str(int(_lowercase ) * int(_lowercase ) )					,		n[i : i + 13] ) )
         for i in range(len(_lowercase ) - 12 ) )


if __name__ == "__main__":
     print(F"""{solution() = }""")