




import argparse
import logging
import os
import re

import tensorflow as tf

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


_lowerCAmelCase						:   Union[str, Any]		=   logging.getLogger(__name__)

_lowerCAmelCase						:   Optional[int]		=   tf.data.AUTOTUNE



def 						lowerCAmelCase		(		):




 """simple docstring"""







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

 UpperCAmelCase__							 =						parser.parse_args()
 return args



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]	):




 """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(_lowerCAmelCase	)
 tf.tpu.experimental.initialize_tpu_system(_lowerCAmelCase	)

 return tpu



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]	):




 """simple docstring"""







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

 return num_samples



def 						lowerCAmelCase		(		_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   Optional[int]						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   List[Any]=None	):




 """simple docstring"""







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



def 						lowerCAmelCase		(		_lowerCAmelCase							:   str	):




 """simple docstring"""







 if not args.no_tpu:
  UpperCAmelCase__							 =						initialize_tpu(_lowerCAmelCase	)
  UpperCAmelCase__							 =						tf.distribute.TPUStrategy(_lowerCAmelCase	)
 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(_lowerCAmelCase	)

 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(_lowerCAmelCase	)
  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=_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							:   Tuple	):
  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(_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.
 UpperCAmelCase__							 =						DataCollatorForLanguageModeling(
     tokenizer=_lowerCAmelCase						,							mlm_probability=args.mlm_probability						,							mlm=_lowerCAmelCase						,							return_tensors="tf"	)

 def mask_with_collator(_lowerCAmelCase							:   Any	):
  # 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(_lowerCAmelCase	)						,							mask_token_id=tokenizer.mask_token_id						,							special_tokens_mask=_lowerCAmelCase						,							)
  return batch

 UpperCAmelCase__							 =						args.per_replica_batch_size * strategy.num_replicas_in_sync

 UpperCAmelCase__							 =						prepare_dataset(
     _lowerCAmelCase						,							decode_fn=_lowerCAmelCase						,							mask_fn=_lowerCAmelCase						,							batch_size=_lowerCAmelCase						,							shuffle=_lowerCAmelCase						,							shuffle_buffer_size=args.shuffle_buffer_size						,							)

 UpperCAmelCase__							 =						prepare_dataset(
     _lowerCAmelCase						,							decode_fn=_lowerCAmelCase						,							mask_fn=_lowerCAmelCase						,							batch_size=_lowerCAmelCase						,							shuffle=_lowerCAmelCase						,							)

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











from typing import TYPE_CHECKING

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


_lowerCAmelCase						:   Union[str, Any]		=   {
    "configuration_bert": ["BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "BertConfig", "BertOnnxConfig"],
    "tokenization_bert": ["BasicTokenizer", "BertTokenizer", "WordpieceTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["BertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   int		=   [
      "BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "BertForMaskedLM",
      "BertForMultipleChoice",
      "BertForNextSentencePrediction",
      "BertForPreTraining",
      "BertForQuestionAnswering",
      "BertForSequenceClassification",
      "BertForTokenClassification",
      "BertLayer",
      "BertLMHeadModel",
      "BertModel",
      "BertPreTrainedModel",
      "load_tf_weights_in_bert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Dict		=   [
      "TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFBertEmbeddings",
      "TFBertForMaskedLM",
      "TFBertForMultipleChoice",
      "TFBertForNextSentencePrediction",
      "TFBertForPreTraining",
      "TFBertForQuestionAnswering",
      "TFBertForSequenceClassification",
      "TFBertForTokenClassification",
      "TFBertLMHeadModel",
      "TFBertMainLayer",
      "TFBertModel",
      "TFBertPreTrainedModel",
  ]
try:
  if not is_tensorflow_text_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["TFBertTokenizer"]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Optional[Any]		=   [
      "FlaxBertForCausalLM",
      "FlaxBertForMaskedLM",
      "FlaxBertForMultipleChoice",
      "FlaxBertForNextSentencePrediction",
      "FlaxBertForPreTraining",
      "FlaxBertForQuestionAnswering",
      "FlaxBertForSequenceClassification",
      "FlaxBertForTokenClassification",
      "FlaxBertModel",
      "FlaxBertPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig, BertOnnxConfig
  from .tokenization_bert import BasicTokenizer, BertTokenizer, WordpieceTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_fast import BertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_bert import (
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BertForMaskedLM,
        BertForMultipleChoice,
        BertForNextSentencePrediction,
        BertForPreTraining,
        BertForQuestionAnswering,
        BertForSequenceClassification,
        BertForTokenClassification,
        BertLayer,
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_bert import (
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFBertEmbeddings,
        TFBertForMaskedLM,
        TFBertForMultipleChoice,
        TFBertForNextSentencePrediction,
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
        TFBertLMHeadModel,
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
    )

  try:
    if not is_tensorflow_text_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_tf import TFBertTokenizer

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_bert import (
        FlaxBertForCausalLM,
        FlaxBertForMaskedLM,
        FlaxBertForMultipleChoice,
        FlaxBertForNextSentencePrediction,
        FlaxBertForPreTraining,
        FlaxBertForQuestionAnswering,
        FlaxBertForSequenceClassification,
        FlaxBertForTokenClassification,
        FlaxBertModel,
        FlaxBertPreTrainedModel,
    )

else:
  import sys

  _lowerCAmelCase						:   Union[str, Any]		=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





