




import argparse
import logging
import os
import re

import tensorflow as tf

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


lowerCamelCase__			:       Optional[Any]         =							logging.getLogger(__name__)

lowerCamelCase__			:       Union[str, Any]         =							tf.data.AUTOTUNE

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

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

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

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

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

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

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

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

      SCREAMING_SNAKE_CASE_     						=      parser.parse_args()
      return args

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

      return tpu

def  UpperCAmelCase_							(				__UpperCAmelCase						:     List[Any]							) ->			List[Any]:
      SCREAMING_SNAKE_CASE_     						=      0
      for file in file_list:
            SCREAMING_SNAKE_CASE_     						=      file.split('/'							)[-1]
            SCREAMING_SNAKE_CASE_     						=      re.search(r'-\d+-(\d+)\.tfrecord'      ,      __UpperCAmelCase							).group(1							)
            SCREAMING_SNAKE_CASE_     						=      int(__UpperCAmelCase							)
            num_samples += sample_count

      return num_samples

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

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[Any]							) ->			Dict:
      if not args.no_tpu:
            SCREAMING_SNAKE_CASE_     						=      initialize_tpu(__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_     						=      tf.distribute.TPUStrategy(__UpperCAmelCase							)
      else:
            SCREAMING_SNAKE_CASE_     						=      tf.distribute.OneDeviceStrategy(device='/gpu:0'							)

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

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

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

      SCREAMING_SNAKE_CASE_     						=      count_samples(__UpperCAmelCase							)

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

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

      def decode_fn(__UpperCAmelCase						:     str							):
            SCREAMING_SNAKE_CASE_     						=      {
                '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(__UpperCAmelCase      ,      __UpperCAmelCase							)

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

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

      SCREAMING_SNAKE_CASE_     						=      args.per_replica_batch_size * strategy.num_replicas_in_sync

      SCREAMING_SNAKE_CASE_     						=      prepare_dataset(
          __UpperCAmelCase      ,      decode_fn=__UpperCAmelCase      ,      mask_fn=__UpperCAmelCase      ,      batch_size=__UpperCAmelCase      ,      shuffle=__UpperCAmelCase      ,      shuffle_buffer_size=args.shuffle_buffer_size      ,      )

      SCREAMING_SNAKE_CASE_     						=      prepare_dataset(
          __UpperCAmelCase      ,      decode_fn=__UpperCAmelCase      ,      mask_fn=__UpperCAmelCase      ,      batch_size=__UpperCAmelCase      ,      shuffle=__UpperCAmelCase      ,      )

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

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

      model.save_pretrained(args.output_dir							)


if __name__ == "__main__":
   lowerCamelCase__			:       Tuple         =							parse_args()
   main(args)





import gc
import threading
import time

import psutil
import torch





class 	lowerCamelCase_							:


 '''simple docstring'''







 def __init__(							self  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_     						=      psutil.Process()
       SCREAMING_SNAKE_CASE_     						=      False



 def 				lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_     						=      -1

       while True:
             SCREAMING_SNAKE_CASE_     						=      max(self.process.memory_info().rss    ,    self.cpu_memory_peak						)

             # can't sleep or will not catch the peak right (this comment is here on purpose)
             if not self.peak_monitoring:
                   break



 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      True
       SCREAMING_SNAKE_CASE_     						=      threading.Thread(target=self.peak_monitor						)
       SCREAMING_SNAKE_CASE_     						=      True
       self.thread.start()




 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      False
       self.thread.join()
       return self.cpu_memory_peak


lowerCamelCase__			:       List[str]         =							PeakCPUMemory()

def  UpperCAmelCase_							(				) ->			Tuple:
      # Time
      SCREAMING_SNAKE_CASE_     						=      {'time': time.time()}

      gc.collect()
      torch.cuda.empty_cache()

      # CPU mem
      SCREAMING_SNAKE_CASE_     						=      psutil.Process().memory_info().rss
      cpu_peak_tracker.start()

      # GPU mem
      for i in range(torch.cuda.device_count()							):
            SCREAMING_SNAKE_CASE_     						=      torch.cuda.memory_allocated(__UpperCAmelCase							)
      torch.cuda.reset_peak_memory_stats()

      return measures

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			Optional[Any]:
      # Time
      SCREAMING_SNAKE_CASE_     						=      {'time': time.time() - start_measures['time']}

      gc.collect()
      torch.cuda.empty_cache()

      # CPU mem
      SCREAMING_SNAKE_CASE_     						=      (psutil.Process().memory_info().rss - start_measures['cpu']) / 2**20
      SCREAMING_SNAKE_CASE_     						=      (cpu_peak_tracker.stop() - start_measures['cpu']) / 2**20

      # GPU mem
      for i in range(torch.cuda.device_count()							):
            SCREAMING_SNAKE_CASE_     						=      (torch.cuda.memory_allocated(__UpperCAmelCase							) - start_measures[str(__UpperCAmelCase							)]) / 2**20
            SCREAMING_SNAKE_CASE_     						=      (torch.cuda.max_memory_allocated(__UpperCAmelCase							) - start_measures[str(__UpperCAmelCase							)]) / 2**20

      return measures





def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]      ,      __UpperCAmelCase						:     Any							) ->			Optional[Any]:
      print(f"{description}:"							)
      print(f"- Time: {measures['time']:.2f}s"							)
      for i in range(torch.cuda.device_count()							):
            print(f"- GPU {i} allocated: {measures[str(__UpperCAmelCase							)]:.2f}MiB"							)
            SCREAMING_SNAKE_CASE_     						=      measures[f"{i}-peak"]
            print(f"- GPU {i} peak: {peak:.2f}MiB"							)
      print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB"							)
      print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB"							)