



"""simple docstring"""





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'''






						snake_case_   : List[str]						=			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**1_8					, 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_1_2					, 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.15					, 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."""    )

						snake_case_   : List[Any]						=			parser.parse_args()
						return args
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Any    ):


						'''simple docstring'''






						try:
												if args.tpu_name:
																		snake_case_   : Optional[int]						=			tf.distribute.cluster_resolver.TPUClusterResolver(
																		    args.tpu_name					, zone=args.tpu_zone					, project=args.gcp_project    )
												else:
																		snake_case_   : 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(__UpperCamelCase    )
						tf.tpu.experimental.initialize_tpu_system(__UpperCamelCase    )

						return tpu
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]    ):


						'''simple docstring'''






						snake_case_   : Optional[int]						=			0
						for file in file_list:
												snake_case_   : List[str]						=			file.split("""/"""    )[-1]
												snake_case_   : Tuple						=			re.search(r"""-\d+-(\d+)\.tfrecord"""					, __UpperCamelCase    ).group(1    )
												snake_case_   : int						=			int(__UpperCamelCase    )
												num_samples += sample_count

						return num_samples
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[Any]					, __UpperCamelCase      :			Tuple					, __UpperCamelCase      :			List[str]					, __UpperCamelCase      :			int					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			Any=None    ):


						'''simple docstring'''






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


						'''simple docstring'''






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

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

						snake_case_   : Union[str, Any]						=			AutoTokenizer.from_pretrained(args.tokenizer    )
						snake_case_   : Optional[int]						=			AutoConfig.from_pretrained(args.pretrained_model_config    )
						snake_case_   : int						=			tokenizer.vocab_size

						snake_case_   : List[str]						=			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}.'    )
						snake_case_   : Tuple						=			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}.'    )

						snake_case_   : Optional[int]						=			count_samples(__UpperCamelCase    )

						snake_case_   : Optional[int]						=			num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
						snake_case_   : Dict						=			steps_per_epoch * args.num_epochs

						with strategy.scope():
												snake_case_   : List[str]						=			TFAutoModelForMaskedLM.from_config(__UpperCamelCase    )
												model(model.dummy_inputs    )  # Pass some dummy inputs through the model to ensure all the weights are built
												snake_case_			,		snake_case_   : Union[str, Any]						=			create_optimizer(
												    num_train_steps=__UpperCamelCase					, num_warmup_steps=total_train_steps // 2_0					, 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      :			Union[str, Any]    ):
												snake_case_   : Tuple						=			{
												    """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.
						snake_case_   : int						=			DataCollatorForLanguageModeling(
						    tokenizer=__UpperCamelCase					, mlm_probability=args.mlm_probability					, mlm=__UpperCamelCase					, return_tensors="""tf"""    )

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

						snake_case_   : Optional[int]						=			args.per_replica_batch_size * strategy.num_replicas_in_sync

						snake_case_   : Any						=			prepare_dataset(
						    __UpperCamelCase					, decode_fn=__UpperCamelCase					, mask_fn=__UpperCamelCase					, batch_size=__UpperCamelCase					, shuffle=__UpperCamelCase					, shuffle_buffer_size=args.shuffle_buffer_size					, )

						snake_case_   : List[Any]						=			prepare_dataset(
						    __UpperCamelCase					, decode_fn=__UpperCamelCase					, mask_fn=__UpperCamelCase					, batch_size=__UpperCamelCase					, shuffle=__UpperCamelCase					, )

						snake_case_   : List[str]						=			[]
						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       :					int      					=					parse_args()
					main(args)






"""simple docstring"""





import unittest

import numpy as np

from transformers import RobertaPreLayerNormConfig, 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():
					import jax.numpy as jnp

					from transformers.models.roberta_prelayernorm.modeling_flax_roberta_prelayernorm import (
					    FlaxRobertaPreLayerNormForCausalLM,
					    FlaxRobertaPreLayerNormForMaskedLM,
					    FlaxRobertaPreLayerNormForMultipleChoice,
					    FlaxRobertaPreLayerNormForQuestionAnswering,
					    FlaxRobertaPreLayerNormForSequenceClassification,
					    FlaxRobertaPreLayerNormForTokenClassification,
					    FlaxRobertaPreLayerNormModel,
					)



class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		def __init__(					self							,    _lowercase							,    _lowercase=1_3							,    _lowercase=7							,    _lowercase=True							,    _lowercase=True							,    _lowercase=True							,    _lowercase=True							,    _lowercase=9_9							,    _lowercase=3_2							,    _lowercase=5							,    _lowercase=4							,    _lowercase=3_7							,    _lowercase="gelu"							,    _lowercase=0.1							,    _lowercase=0.1							,    _lowercase=5_1_2							,    _lowercase=1_6							,    _lowercase=2							,    _lowercase=0.02							,    _lowercase=4							,    ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Optional[int]						=			parent
								snake_case_   : Dict						=			batch_size
								snake_case_   : Any						=			seq_length
								snake_case_   : Tuple						=			is_training
								snake_case_   : Dict						=			use_attention_mask
								snake_case_   : int						=			use_token_type_ids
								snake_case_   : List[Any]						=			use_labels
								snake_case_   : List[str]						=			vocab_size
								snake_case_   : str						=			hidden_size
								snake_case_   : Optional[Any]						=			num_hidden_layers
								snake_case_   : List[Any]						=			num_attention_heads
								snake_case_   : Any						=			intermediate_size
								snake_case_   : Optional[Any]						=			hidden_act
								snake_case_   : Tuple						=			hidden_dropout_prob
								snake_case_   : List[Any]						=			attention_probs_dropout_prob
								snake_case_   : int						=			max_position_embeddings
								snake_case_   : Dict						=			type_vocab_size
								snake_case_   : Dict						=			type_sequence_label_size
								snake_case_   : Optional[Any]						=			initializer_range
								snake_case_   : Tuple						=			num_choices





		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : int						=			ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size   )

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

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

								snake_case_   : Tuple						=			RobertaPreLayerNormConfig(
								    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=_lowercase							,    initializer_range=self.initializer_range							,    )

								return config, input_ids, token_type_ids, attention_mask





		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								snake_case_   : Any						=			self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_   : List[Any]						=			config_and_inputs
								snake_case_   : int						=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
								return config, inputs_dict





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Tuple						=			self.prepare_config_and_inputs()
								snake_case_			,		snake_case_			,		snake_case_			,		snake_case_   : Any						=			config_and_inputs

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

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




@require_flax
# Copied from tests.models.roberta.test_modelling_flax_roberta.FlaxRobertaPreLayerNormModelTest with ROBERTA->ROBERTA_PRELAYERNORM,Roberta->RobertaPreLayerNorm,roberta-base->andreasmadsen/efficient_mlm_m0.40
class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       True

		_lowerCamelCase			  =       (
		    (
		        FlaxRobertaPreLayerNormModel,
		        FlaxRobertaPreLayerNormForCausalLM,
		        FlaxRobertaPreLayerNormForMaskedLM,
		        FlaxRobertaPreLayerNormForSequenceClassification,
		        FlaxRobertaPreLayerNormForTokenClassification,
		        FlaxRobertaPreLayerNormForMultipleChoice,
		        FlaxRobertaPreLayerNormForQuestionAnswering,
		    )
		    if is_flax_available()
		    else ()
		)





		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			FlaxRobertaPreLayerNormModelTester(self   )





		@slow
		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								for model_class_name in self.all_model_classes:
														snake_case_   : Tuple						=			model_class_name.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )
														snake_case_   : Tuple						=			model(np.ones((1, 1)   )   )
														self.assertIsNotNone(_lowercase   )




@require_flax
class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		@slow
		def    UpperCAmelCase__							(					self   ) ->						Union[str, Any]:







								'''simple docstring'''

								snake_case_   : Union[str, Any]						=			FlaxRobertaPreLayerNormForMaskedLM.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )

								snake_case_   : List[str]						=			np.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]							,    dtype=jnp.intaa   )
								snake_case_   : str						=			model(_lowercase   )[0]
								snake_case_   : int						=			[1, 1_1, 5_0_2_6_5]
								self.assertEqual(list(output.shape   )							,    _lowercase   )
								# compare the actual values for a slice.
								snake_case_   : Tuple						=			np.array(
								    [[[40.4880, 18.0199, -5.2367], [-1.8877, -4.0885, 10.7085], [-2.2613, -5.6110, 7.2665]]]							,    dtype=np.floataa   )
								self.assertTrue(np.allclose(output[:, :3, :3]							,    _lowercase							,    atol=1E-4   )   )





		@slow
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Any						=			FlaxRobertaPreLayerNormModel.from_pretrained("""andreasmadsen/efficient_mlm_m0.40"""							,    from_pt=_lowercase   )

								snake_case_   : Dict						=			np.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]							,    dtype=jnp.intaa   )
								snake_case_   : Any						=			model(_lowercase   )[0]
								# compare the actual values for a slice.
								snake_case_   : Optional[Any]						=			np.array(
								    [[[0.0208, -0.0356, 0.0237], [-0.1569, -0.0411, -0.2626], [0.1879, 0.0125, -0.0089]]]							,    dtype=np.floataa   )
								self.assertTrue(np.allclose(output[:, :3, :3]							,    _lowercase							,    atol=1E-4   )   )

