
"""simple docstring"""





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 	__a     (		):
   UpperCAmelCase_					:					Any			     =    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_					:					Dict			     =    parser.parse_args()
   return args







def 	__a     (		__lowerCamelCase  ):
   try:
      if args.tpu_name:
         UpperCAmelCase_					:					Optional[Any]			     =    tf.distribute.cluster_resolver.TPUClusterResolver(
             args.tpu_name,		zone=args.tpu_zone,		project=args.gcp_project  )
      else:
         UpperCAmelCase_					:					Dict			     =    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 	__a     (		__lowerCamelCase  ):
   UpperCAmelCase_					:					Optional[int]			     =    0
   for file in file_list:
      UpperCAmelCase_					:					Union[str, Any]			     =    file.split("/"  )[-1]
      UpperCAmelCase_					:					Union[str, Any]			     =    re.search(r"-\d+-(\d+)\.tfrecord",		__lowerCamelCase  ).group(1  )
      UpperCAmelCase_					:					List[Any]			     =    int(__lowerCamelCase  )
      num_samples += sample_count

   return num_samples







def 	__a     (		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase=None  ):
   UpperCAmelCase_					:					Optional[int]			     =    count_samples(__lowerCamelCase  )
   UpperCAmelCase_					:					Optional[Any]			     =    tf.data.Dataset.from_tensor_slices(__lowerCamelCase  )
   if shuffle:
      UpperCAmelCase_					:					Dict			     =    dataset.shuffle(len(__lowerCamelCase  )  )
   UpperCAmelCase_					:					Any			     =    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_					:					str			     =    dataset.apply(tf.data.experimental.assert_cardinality(__lowerCamelCase  )  )
   UpperCAmelCase_					:					Any			     =    dataset.map(__lowerCamelCase,		num_parallel_calls=__lowerCamelCase  )
   if shuffle:
      assert shuffle_buffer_size is not None
      UpperCAmelCase_					:					Optional[int]			     =    dataset.shuffle(args.shuffle_buffer_size  )
   UpperCAmelCase_					:					Tuple			     =    dataset.batch(__lowerCamelCase,		drop_remainder=__lowerCamelCase  )
   UpperCAmelCase_					:					Tuple			     =    dataset.map(__lowerCamelCase,		num_parallel_calls=__lowerCamelCase  )
   UpperCAmelCase_					:					Optional[Any]			     =    dataset.prefetch(__lowerCamelCase  )
   return dataset







def 	__a     (		__lowerCamelCase  ):
   if not args.no_tpu:
      UpperCAmelCase_					:					List[str]			     =    initialize_tpu(__lowerCamelCase  )
      UpperCAmelCase_					:					Union[str, Any]			     =    tf.distribute.TPUStrategy(__lowerCamelCase  )
   else:
      UpperCAmelCase_					:					Dict			     =    tf.distribute.OneDeviceStrategy(device="/gpu:0"  )

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

   UpperCAmelCase_					:					Dict			     =    AutoTokenizer.from_pretrained(args.tokenizer  )
   UpperCAmelCase_					:					Optional[Any]			     =    AutoConfig.from_pretrained(args.pretrained_model_config  )
   UpperCAmelCase_					:					Any			     =    tokenizer.vocab_size

   UpperCAmelCase_					:					Dict			     =    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_					:					Union[str, Any]			     =    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_					:					Union[str, Any]			     =    count_samples(__lowerCamelCase  )

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

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

   def mask_with_collator(__lowerCamelCase  ):
      # TF really needs an isin() function
      UpperCAmelCase_					:					List[Any]			     =    (
          ~tf.cast(batch["attention_mask"],		tf.bool  )
          | (batch["input_ids"] == tokenizer.cls_token_id)
          | (batch["input_ids"] == tokenizer.sep_token_id)
      )
      UpperCAmelCase_     ,   UpperCAmelCase_					:					Dict			     =    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_					:					List[str]			     =    args.per_replica_batch_size * strategy.num_replicas_in_sync

   UpperCAmelCase_					:					List[str]			     =    prepare_dataset(
       __lowerCamelCase,		decode_fn=__lowerCamelCase,		mask_fn=__lowerCamelCase,		batch_size=__lowerCamelCase,		shuffle=__lowerCamelCase,		shuffle_buffer_size=args.shuffle_buffer_size,		)

   UpperCAmelCase_					:					int			     =    prepare_dataset(
       __lowerCamelCase,		decode_fn=__lowerCamelCase,		mask_fn=__lowerCamelCase,		batch_size=__lowerCamelCase,		shuffle=__lowerCamelCase,		)

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







'''simple docstring'''




import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
							from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
							from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
							from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


a   :			str					  =    DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

a   :			int					  =    """main"""
# Default branch name
a   :			Any					  =    """f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"""
# One particular commit (not the top of `main`)
a   :			str					  =    """aaaaaaa"""
# This commit does not exist, so we should 404.

a   :			int					  =    """d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"""
# Sha-1 of config.json on the top of `main`, for checking purposes
a   :			Any					  =    """4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"""



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			List[str]:
							print("""Welcome!"""		)
							yield
							print("""Bye!"""		)



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			Optional[int]:
							print("""Bonjour!"""		)
							yield
							print("""Au revoir!"""		)

class 			UpperCamelCase_ (							unittest.TestCase			):





	def  _lowercase(  self    )			->			List[Any]:
								# If the spec is missing, importlib would not be able to import the module dynamically.
								assert transformers.__spec__ is not None
								assert importlib.util.find_spec("""transformers"""    ) is not None




class 			UpperCamelCase_ (							unittest.TestCase			):





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Tuple:
								with ContextManagers([]    ):
															print("""Transformers are awesome!"""    )
								# The print statement adds a new line at the end of the output
								self.assertEqual(mock_stdout.getvalue()		,		"""Transformers are awesome!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Dict:
								with ContextManagers([context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Welcome!\nTransformers are awesome!\nBye!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Union[str, Any]:
								with ContextManagers([context_fr(), context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English and French welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"""    )





	@require_torch
	def  _lowercase(  self    )			->			Optional[int]:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )





	@require_tf
	def  _lowercase(  self    )			->			int:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )






	@require_flax
	def  _lowercase(  self    )			->			Any:
								# Flax models don't have labels
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		[]    )

