


'''simple docstring'''
import argparse
import logging
import os
import re

import tensorflow as tf

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


__snake_case											=					logging.getLogger(__name__)

__snake_case											=					tf.data.AUTOTUNE






def       a				(      )  ->     int:



					'''simple docstring'''






					UpperCamelCase__       :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**18			, 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=512			, 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.1_5			, 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.'''       )

					UpperCamelCase__       :int							 =			parser.parse_args()
					return args






def       a				(      __a       )  ->     Any:



					'''simple docstring'''






					try:
										if args.tpu_name:
															UpperCamelCase__       :int							 =			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(_UpperCAmelCase       )
					tf.tpu.experimental.initialize_tpu_system(_UpperCAmelCase       )

					return tpu






def       a				(      __a       )  ->     Any:



					'''simple docstring'''






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

					return num_samples






def       a				(      __a			, __a			, __a			, __a			, __a			, __a=None       )  ->     List[str]:



					'''simple docstring'''






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






def       a				(      __a       )  ->     str:



					'''simple docstring'''






					if not args.no_tpu:
										UpperCamelCase__       :Optional[Any]							 =			initialize_tpu(_UpperCAmelCase       )
										UpperCamelCase__       :List[Any]							 =			tf.distribute.TPUStrategy(_UpperCAmelCase       )
					else:
										UpperCamelCase__       :str							 =			tf.distribute.OneDeviceStrategy(device='''/gpu:0'''       )

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

					UpperCamelCase__       :List[Any]							 =			AutoTokenizer.from_pretrained(args.tokenizer       )
					UpperCamelCase__       :str							 =			AutoConfig.from_pretrained(args.pretrained_model_config       )
					UpperCamelCase__       :int							 =			tokenizer.vocab_size

					UpperCamelCase__       :Tuple							 =			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__       :List[str]							 =			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(_UpperCAmelCase       )

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

					with strategy.scope():
										UpperCamelCase__       :int							 =			TFAutoModelForMaskedLM.from_config(_UpperCAmelCase       )
										model(model.dummy_inputs       )  # Pass some dummy inputs through the model to ensure all the weights are built
										UpperCamelCase__       :int							 =			create_optimizer(
										    num_train_steps=_UpperCAmelCase			, 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=_UpperCAmelCase			, metrics=['''accuracy''']       )

					def decode_fn(__a       ):
										UpperCamelCase__       :Union[str, Any]							 =			{
										    '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.
					UpperCamelCase__       :int							 =			DataCollatorForLanguageModeling(
					    tokenizer=_UpperCAmelCase			, mlm_probability=args.mlm_probability			, mlm=_UpperCAmelCase			, return_tensors='''tf'''       )

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

					UpperCamelCase__       :Tuple							 =			args.per_replica_batch_size * strategy.num_replicas_in_sync

					UpperCamelCase__       :str							 =			prepare_dataset(
					    _UpperCAmelCase			, decode_fn=_UpperCAmelCase			, mask_fn=_UpperCAmelCase			, batch_size=_UpperCAmelCase			, shuffle=_UpperCAmelCase			, shuffle_buffer_size=args.shuffle_buffer_size			, )

					UpperCamelCase__       :List[str]							 =			prepare_dataset(
					    _UpperCAmelCase			, decode_fn=_UpperCAmelCase			, mask_fn=_UpperCAmelCase			, batch_size=_UpperCAmelCase			, shuffle=_UpperCAmelCase			, )

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



'''simple docstring'''
import os
import unittest

from huggingface_hub.utils import are_progress_bars_disabled

import transformers.models.bart.tokenization_bart
from transformers import logging
from transformers.testing_utils import CaptureLogger, mockenv, mockenv_context
from transformers.utils.logging import disable_progress_bar, enable_progress_bar
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			logging.get_logger()

							# the current default level is logging.WARNING
							UpperCamelCase__       :List[Any]							 =			logging.get_verbosity()

							logging.set_verbosity_error()
							self.assertEqual(logger.getEffectiveLevel()   ,       logging.get_verbosity()							)

							logging.set_verbosity_warning()
							self.assertEqual(logger.getEffectiveLevel()   ,       logging.get_verbosity()							)

							logging.set_verbosity_info()
							self.assertEqual(logger.getEffectiveLevel()   ,       logging.get_verbosity()							)

							logging.set_verbosity_debug()
							self.assertEqual(logger.getEffectiveLevel()   ,       logging.get_verbosity()							)

							# restore to the original level
							logging.set_verbosity(UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			logging.get_verbosity()

							UpperCamelCase__       :Union[str, Any]							 =			logging.get_logger('''transformers.models.bart.tokenization_bart'''							)
							UpperCamelCase__       :Optional[Any]							 =			'''Testing 1, 2, 3'''

							# should be able to log warnings (if default settings weren't overridden by `pytest --log-level-all`)
							if level_origin <= logging.WARNING:
												with CaptureLogger(UpperCamelCase_							) as cl:
																	logger.warning(UpperCamelCase_							)
												self.assertEqual(cl.out   ,       msg + '''\n'''							)

							# this is setting the level for all of `transformers.*` loggers
							logging.set_verbosity_error()

							# should not be able to log warnings
							with CaptureLogger(UpperCamelCase_							) as cl:
												logger.warning(UpperCamelCase_							)
							self.assertEqual(cl.out   ,       ''''''							)

							# should be able to log warnings again
							logging.set_verbosity_warning()
							with CaptureLogger(UpperCamelCase_							) as cl:
												logger.warning(UpperCamelCase_							)
							self.assertEqual(cl.out   ,       msg + '''\n'''							)

							# restore to the original level
							logging.set_verbosity(UpperCamelCase_							)




		@mockenv(TRANSFORMERS_VERBOSITY='''error'''							)
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							transformers.utils.logging._reset_library_root_logger()
							# this action activates the env var
							UpperCamelCase__       :Tuple							 =			logging.get_logger('''transformers.models.bart.tokenization_bart'''							)

							UpperCamelCase__       :Any							 =			os.getenv('''TRANSFORMERS_VERBOSITY'''   ,       UpperCamelCase_							)
							UpperCamelCase__       :Optional[Any]							 =			logging.log_levels[env_level_str]

							UpperCamelCase__       :int							 =			logging.get_verbosity()
							self.assertEqual(
							    UpperCamelCase_   ,       UpperCamelCase_   ,       F'''TRANSFORMERS_VERBOSITY={env_level_str}/{env_level}, but internal verbosity is {current_level}'''   ,       )

							# restore to the original level
							UpperCamelCase__       :Union[str, Any]							 =			''''''
							transformers.utils.logging._reset_library_root_logger()




		@mockenv(TRANSFORMERS_VERBOSITY='''super-error'''							)
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							transformers.utils.logging._reset_library_root_logger()
							UpperCamelCase__       :Dict							 =			logging.logging.getLogger()
							with CaptureLogger(UpperCamelCase_							) as cl:
												# this action activates the env var
												logging.get_logger('''transformers.models.bart.tokenization_bart'''							)
							self.assertIn('''Unknown option TRANSFORMERS_VERBOSITY=super-error'''   ,       cl.out							)

							# no need to restore as nothing was changed




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							transformers.utils.logging._reset_library_root_logger()

							UpperCamelCase__       :Optional[int]							 =			logging.get_logger('''transformers.models.bart.tokenization_bart'''							)
							UpperCamelCase__       :int							 =			'''Testing 1, 2, 3'''

							with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS='''1'''							):
												# nothing should be logged as env var disables this method
												with CaptureLogger(UpperCamelCase_							) as cl:
																	logger.warning_advice(UpperCamelCase_							)
												self.assertEqual(cl.out   ,       ''''''							)

							with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS=''''''							):
												# should log normally as TRANSFORMERS_NO_ADVISORY_WARNINGS is unset
												with CaptureLogger(UpperCamelCase_							) as cl:
																	logger.warning_advice(UpperCamelCase_							)
												self.assertEqual(cl.out   ,       msg + '''\n'''							)



def       a				(      )  ->     str:



					'''simple docstring'''






					disable_progress_bar()
					assert are_progress_bars_disabled()

					enable_progress_bar()
					assert not are_progress_bars_disabled()