
import argparse
import logging
import os
import re

import tensorflow as tf

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


lowerCamelCase_							=     logging.getLogger(__name__)

lowerCamelCase_							=     tf.data.AUTOTUNE






def 						__magic_name__ (    ):



     '''simple docstring'''




     UpperCamelCase__				=				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__				=				parser.parse_args()
     return args






def 						__magic_name__ (    __a			:       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(_lowercase						)
     tf.tpu.experimental.initialize_tpu_system(_lowercase						)

     return tpu






def 						__magic_name__ (    __a			:       int						):



     '''simple docstring'''




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

     return num_samples






def 						__magic_name__ (    __a			:       List[Any]       ,		__a			:       Optional[Any]       ,		__a			:       Any       ,		__a			:       List[str]       ,		__a			:       Any       ,		__a			:       int=None						):



     '''simple docstring'''




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






def 						__magic_name__ (    __a			:       List[str]						):



     '''simple docstring'''




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

     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(_lowercase						)
          model(model.dummy_inputs						)  # Pass some dummy inputs through the model to ensure all the weights are built
          UpperCamelCase__				=				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(__a			:       Dict						):
          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(_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__				=				DataCollatorForLanguageModeling(
         tokenizer=_lowercase       ,		mlm_probability=args.mlm_probability       ,		mlm=_lowercase       ,		return_tensors="""tf"""						)

     def mask_with_collator(__a			:       List[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__				=				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__				=				args.per_replica_batch_size * strategy.num_replicas_in_sync

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

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

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










"""simple docstring"""
import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
						from transformers.models.bert.modeling_flax_bert import (
						    FlaxBertForMaskedLM,
						    FlaxBertForMultipleChoice,
						    FlaxBertForNextSentencePrediction,
						    FlaxBertForPreTraining,
						    FlaxBertForQuestionAnswering,
						    FlaxBertForSequenceClassification,
						    FlaxBertForTokenClassification,
						    FlaxBertModel,
						)
class 				lowerCamelCase  (			unittest.TestCase				):






					'''simple docstring'''





					def __init__(       self:							List[Any]						,     snake_case:							List[str]						,     snake_case:							Optional[Any]=13						,     snake_case:							List[str]=7						,     snake_case:							Dict=True						,     snake_case:							List[str]=True						,     snake_case:							Optional[int]=True						,     snake_case:							Any=True						,     snake_case:							Optional[Any]=99						,     snake_case:							Tuple=32						,     snake_case:							Tuple=5						,     snake_case:							Dict=4						,     snake_case:							Optional[Any]=37						,     snake_case:							Union[str, Any]="gelu"						,     snake_case:							Tuple=0.1						,     snake_case:							List[Any]=0.1						,     snake_case:							List[str]=512						,     snake_case:							Optional[int]=16						,     snake_case:							int=2						,     snake_case:							List[Any]=0.0_2						,     snake_case:							Union[str, Any]=4						,     )   ->	List[str]:
											snake_case_ :Dict						=			parent
											snake_case_ :Any						=			batch_size
											snake_case_ :Any						=			seq_length
											snake_case_ :List[str]						=			is_training
											snake_case_ :Optional[Any]						=			use_attention_mask
											snake_case_ :Dict						=			use_token_type_ids
											snake_case_ :Union[str, Any]						=			use_labels
											snake_case_ :str						=			vocab_size
											snake_case_ :int						=			hidden_size
											snake_case_ :List[str]						=			num_hidden_layers
											snake_case_ :Dict						=			num_attention_heads
											snake_case_ :Any						=			intermediate_size
											snake_case_ :Tuple						=			hidden_act
											snake_case_ :int						=			hidden_dropout_prob
											snake_case_ :Optional[Any]						=			attention_probs_dropout_prob
											snake_case_ :Any						=			max_position_embeddings
											snake_case_ :Union[str, Any]						=			type_vocab_size
											snake_case_ :Optional[int]						=			type_sequence_label_size
											snake_case_ :Union[str, Any]						=			initializer_range
											snake_case_ :Tuple						=			num_choices






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	str:
											snake_case_ :Optional[Any]						=			ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size						)

											snake_case_ :Union[str, Any]						=			None
											if self.use_attention_mask:
																	snake_case_ :str						=			random_attention_mask([self.batch_size, self.seq_length]						)

											snake_case_ :Any						=			None
											if self.use_token_type_ids:
																	snake_case_ :List[str]						=			ids_tensor([self.batch_size, self.seq_length]						,     self.type_vocab_size						)

											snake_case_ :int						=			BertConfig(
											    vocab_size=self.vocab_size						,     hidden_size=self.hidden_size						,     num_hidden_layers=self.num_hidden_layers						,     num_attention_heads=self.num_attention_heads						,     intermediate_size=self.intermediate_size						,     hidden_act=self.hidden_act						,     hidden_dropout_prob=self.hidden_dropout_prob						,     attention_probs_dropout_prob=self.attention_probs_dropout_prob						,     max_position_embeddings=self.max_position_embeddings						,     type_vocab_size=self.type_vocab_size						,     is_decoder=snake_case						,     initializer_range=self.initializer_range						,     )

											return config, input_ids, token_type_ids, attention_mask






					def    lowerCAmelCase_				(       self:							Optional[int]						)   ->	int:
											snake_case_ :str						=			self.prepare_config_and_inputs()
											snake_case_,       snake_case_,       snake_case_,       snake_case_ :Optional[int]						=			config_and_inputs
											snake_case_ :Union[str, Any]						=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
											return config, inputs_dict






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Any:
											snake_case_ :int						=			self.prepare_config_and_inputs()
											snake_case_,       snake_case_,       snake_case_,       snake_case_ :Dict						=			config_and_inputs

											snake_case_ :Union[str, Any]						=			True
											snake_case_ :Optional[int]						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
											snake_case_ :Tuple						=			ids_tensor([self.batch_size, self.seq_length]						,     vocab_size=2						)

											return (
											    config,
											    input_ids,
											    attention_mask,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)

@require_flax
class 				lowerCamelCase  (			_lowerCAmelCase    ,						unittest.TestCase				):






					'''simple docstring'''





					_A		:			List[str]						=     True

					_A		:			Dict						=     (
					    (
					        FlaxBertModel,
					        FlaxBertForPreTraining,
					        FlaxBertForMaskedLM,
					        FlaxBertForMultipleChoice,
					        FlaxBertForQuestionAnswering,
					        FlaxBertForNextSentencePrediction,
					        FlaxBertForSequenceClassification,
					        FlaxBertForTokenClassification,
					        FlaxBertForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)






					def    lowerCAmelCase_				(       self:							int						)   ->	List[str]:
											snake_case_ :Any						=			FlaxBertModelTester(self						)






					@slow
					def    lowerCAmelCase_				(       self:							List[str]						)   ->	Dict:
											# Only check this for base model, not necessary for all model classes.
											# This will also help speed-up tests.
											snake_case_ :Dict						=			FlaxBertModel.from_pretrained("""bert-base-cased"""						)
											snake_case_ :Dict						=			model(np.ones((1, 1)						)						)
											self.assertIsNotNone(snake_case						)
