





import argparse
import logging
import os
import re

import tensorflow as tf

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


UpperCamelCase__           =			logging.getLogger(__name__)

UpperCamelCase__           =			tf.data.AUTOTUNE





def      _UpperCamelCase  ():






							"""simple docstring"""





							UpperCamelCase__					  =      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.15 ,      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."""							)

							UpperCamelCase__					  =      parser.parse_args()
							return args





def      _UpperCamelCase  (a__ :Optional[int]							):






							"""simple docstring"""





							try:
														if args.tpu_name:
																					UpperCamelCase__					  =      tf.distribute.cluster_resolver.TPUClusterResolver(
																					    args.tpu_name ,      zone=args.tpu_zone ,      project=args.gcp_project							)
														else:
																					UpperCamelCase__					  =      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__ :Tuple							):






							"""simple docstring"""





							UpperCamelCase__					  =      0
							for file in file_list:
														UpperCamelCase__					  =      file.split("""/"""							)[-1]
														UpperCamelCase__					  =      re.search(r"""-\d+-(\d+)\.tfrecord""" ,      a__							).group(1							)
														UpperCamelCase__					  =      int(a__							)
														num_samples += sample_count

							return num_samples





def      _UpperCamelCase  (a__ :str ,      a__ :Tuple ,      a__ :Optional[Any] ,      a__ :List[Any] ,      a__ :Dict ,      a__ :Union[str, Any]=None							):






							"""simple docstring"""





							UpperCamelCase__					  =      count_samples(a__							)
							UpperCamelCase__					  =      tf.data.Dataset.from_tensor_slices(a__							)
							if shuffle:
														UpperCamelCase__					  =      dataset.shuffle(len(a__							)							)
							UpperCamelCase__					  =      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
							UpperCamelCase__					  =      dataset.apply(tf.data.experimental.assert_cardinality(a__							)							)
							UpperCamelCase__					  =      dataset.map(a__ ,      num_parallel_calls=a__							)
							if shuffle:
														assert shuffle_buffer_size is not None
														UpperCamelCase__					  =      dataset.shuffle(args.shuffle_buffer_size							)
							UpperCamelCase__					  =      dataset.batch(a__ ,      drop_remainder=a__							)
							UpperCamelCase__					  =      dataset.map(a__ ,      num_parallel_calls=a__							)
							UpperCamelCase__					  =      dataset.prefetch(a__							)
							return dataset





def      _UpperCamelCase  (a__ :Any							):






							"""simple docstring"""





							if not args.no_tpu:
														UpperCamelCase__					  =      initialize_tpu(a__							)
														UpperCamelCase__					  =      tf.distribute.TPUStrategy(a__							)
							else:
														UpperCamelCase__					  =      tf.distribute.OneDeviceStrategy(device="""/gpu:0"""							)

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

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

							UpperCamelCase__					  =      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__					  =      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__					  =      count_samples(a__							)

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

							with strategy.scope():
														UpperCamelCase__					  =      TFAutoModelForMaskedLM.from_config(a__							)
														model(model.dummy_inputs							)  # Pass some dummy inputs through the model to ensure all the weights are built
														UpperCamelCase__  ,  UpperCamelCase__					  =      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__ :int							):
														UpperCamelCase__					  =      {
														    """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.
							UpperCamelCase__					  =      DataCollatorForLanguageModeling(
							    tokenizer=a__ ,      mlm_probability=args.mlm_probability ,      mlm=a__ ,      return_tensors="""tf"""							)

							def mask_with_collator(a__ :Dict							):
														# TF really needs an isin() function
														UpperCamelCase__					  =      (
														    ~tf.cast(batch["""attention_mask"""] ,      tf.bool							)
														    | (batch["""input_ids"""] == tokenizer.cls_token_id)
														    | (batch["""input_ids"""] == tokenizer.sep_token_id)
														)
														UpperCamelCase__  ,  UpperCamelCase__					  =      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

							UpperCamelCase__					  =      args.per_replica_batch_size * strategy.num_replicas_in_sync

							UpperCamelCase__					  =      prepare_dataset(
							    a__ ,      decode_fn=a__ ,      mask_fn=a__ ,      batch_size=a__ ,      shuffle=a__ ,      shuffle_buffer_size=args.shuffle_buffer_size ,      )

							UpperCamelCase__					  =      prepare_dataset(
							    a__ ,      decode_fn=a__ ,      mask_fn=a__ ,      batch_size=a__ ,      shuffle=a__ ,      )

							UpperCamelCase__					  =      []
							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__":
					UpperCamelCase__           =			parse_args()
					main(args)











"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE__   =  {
    "configuration_roberta_prelayernorm": [
        "ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "RobertaPreLayerNormConfig",
        "RobertaPreLayerNormOnnxConfig",
    ],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "RobertaPreLayerNormForCausalLM",
           "RobertaPreLayerNormForMaskedLM",
           "RobertaPreLayerNormForMultipleChoice",
           "RobertaPreLayerNormForQuestionAnswering",
           "RobertaPreLayerNormForSequenceClassification",
           "RobertaPreLayerNormForTokenClassification",
           "RobertaPreLayerNormModel",
           "RobertaPreLayerNormPreTrainedModel",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFRobertaPreLayerNormForCausalLM",
           "TFRobertaPreLayerNormForMaskedLM",
           "TFRobertaPreLayerNormForMultipleChoice",
           "TFRobertaPreLayerNormForQuestionAnswering",
           "TFRobertaPreLayerNormForSequenceClassification",
           "TFRobertaPreLayerNormForTokenClassification",
           "TFRobertaPreLayerNormMainLayer",
           "TFRobertaPreLayerNormModel",
           "TFRobertaPreLayerNormPreTrainedModel",
       ]

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "FlaxRobertaPreLayerNormForCausalLM",
           "FlaxRobertaPreLayerNormForMaskedLM",
           "FlaxRobertaPreLayerNormForMultipleChoice",
           "FlaxRobertaPreLayerNormForQuestionAnswering",
           "FlaxRobertaPreLayerNormForSequenceClassification",
           "FlaxRobertaPreLayerNormForTokenClassification",
           "FlaxRobertaPreLayerNormModel",
           "FlaxRobertaPreLayerNormPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_roberta_prelayernorm import (
           ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP,
           RobertaPreLayerNormConfig,
           RobertaPreLayerNormOnnxConfig,
       )

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_roberta_prelayernorm import (
                  ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  RobertaPreLayerNormForCausalLM,
                  RobertaPreLayerNormForMaskedLM,
                  RobertaPreLayerNormForMultipleChoice,
                  RobertaPreLayerNormForQuestionAnswering,
                  RobertaPreLayerNormForSequenceClassification,
                  RobertaPreLayerNormForTokenClassification,
                  RobertaPreLayerNormModel,
                  RobertaPreLayerNormPreTrainedModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_roberta_prelayernorm import (
                  TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFRobertaPreLayerNormForCausalLM,
                  TFRobertaPreLayerNormForMaskedLM,
                  TFRobertaPreLayerNormForMultipleChoice,
                  TFRobertaPreLayerNormForQuestionAnswering,
                  TFRobertaPreLayerNormForSequenceClassification,
                  TFRobertaPreLayerNormForTokenClassification,
                  TFRobertaPreLayerNormMainLayer,
                  TFRobertaPreLayerNormModel,
                  TFRobertaPreLayerNormPreTrainedModel,
              )

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_roberta_prelayernorm import (
                  FlaxRobertaPreLayerNormForCausalLM,
                  FlaxRobertaPreLayerNormForMaskedLM,
                  FlaxRobertaPreLayerNormForMultipleChoice,
                  FlaxRobertaPreLayerNormForQuestionAnswering,
                  FlaxRobertaPreLayerNormForSequenceClassification,
                  FlaxRobertaPreLayerNormForTokenClassification,
                  FlaxRobertaPreLayerNormModel,
                  FlaxRobertaPreLayerNormPreTrainedModel,
              )

else:
       import sys

       SCREAMING_SNAKE_CASE__   =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






