


'''simple docstring'''

import argparse
import logging
import os
import re

import tensorflow as tf

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


_UpperCAmelCase				: Dict							    =				logging.getLogger(__name__)

_UpperCAmelCase				: Tuple							    =				tf.data.AUTOTUNE

def 							__magic_name__(						):
						__lowerCAmelCase							       =				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**1_8,				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=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=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.''')

						__lowerCAmelCase							       =				parser.parse_args()
						return args

def 							__magic_name__(						lowerCamelCase):
						try:
												if args.tpu_name:
																		__lowerCAmelCase							       =				tf.distribute.cluster_resolver.TPUClusterResolver(
																		    args.tpu_name,				zone=args.tpu_zone,				project=args.gcp_project)
												else:
																		__lowerCAmelCase							       =				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 							__magic_name__(						lowerCamelCase):
						__lowerCAmelCase							       =				0
						for file in file_list:
												__lowerCAmelCase							       =				file.split('''/''')[-1]
												__lowerCAmelCase							       =				re.search(r'''-\d+-(\d+)\.tfrecord''',				lowerCamelCase).group(1)
												__lowerCAmelCase							       =				int(lowerCamelCase)
												num_samples += sample_count

						return num_samples

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

def 							__magic_name__(						lowerCamelCase):
						if not args.no_tpu:
												__lowerCAmelCase							       =				initialize_tpu(lowerCamelCase)
												__lowerCAmelCase							       =				tf.distribute.TPUStrategy(lowerCamelCase)
						else:
												__lowerCAmelCase							       =				tf.distribute.OneDeviceStrategy(device='''/gpu:0''')

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

						__lowerCAmelCase							       =				AutoTokenizer.from_pretrained(args.tokenizer)
						__lowerCAmelCase							       =				AutoConfig.from_pretrained(args.pretrained_model_config)
						__lowerCAmelCase							       =				tokenizer.vocab_size

						__lowerCAmelCase							       =				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}.""")
						__lowerCAmelCase							       =				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}.""")

						__lowerCAmelCase							       =				count_samples(lowerCamelCase)

						__lowerCAmelCase							       =				num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
						__lowerCAmelCase							       =				steps_per_epoch * args.num_epochs

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

						def decode_fn(lowerCamelCase):
												__lowerCAmelCase							       =				{
												    '''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.
						__lowerCAmelCase							       =				DataCollatorForLanguageModeling(
						    tokenizer=lowerCamelCase,				mlm_probability=args.mlm_probability,				mlm=lowerCamelCase,				return_tensors='''tf''')

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

						__lowerCAmelCase							       =				args.per_replica_batch_size * strategy.num_replicas_in_sync

						__lowerCAmelCase							       =				prepare_dataset(
						    lowerCamelCase,				decode_fn=lowerCamelCase,				mask_fn=lowerCamelCase,				batch_size=lowerCamelCase,				shuffle=lowerCamelCase,				shuffle_buffer_size=args.shuffle_buffer_size,				)

						__lowerCAmelCase							       =				prepare_dataset(
						    lowerCamelCase,				decode_fn=lowerCamelCase,				mask_fn=lowerCamelCase,				batch_size=lowerCamelCase,				shuffle=lowerCamelCase,				)

						__lowerCAmelCase							       =				[]
						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__":
	_UpperCAmelCase				: Tuple							    =				parse_args()
	main(args)





'''simple docstring'''

import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFGPTaLMHeadModel, is_keras_nlp_available, is_tf_available
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import require_keras_nlp, require_tf, slow


if is_tf_available():
	import tensorflow as tf

if is_keras_nlp_available():
	from transformers.models.gpta import TFGPTaTokenizer


_UpperCAmelCase				: str							    =				["""gpt2"""]
_UpperCAmelCase				: Optional[int]							    =				"""gpt2"""

if is_tf_available():


	class       a__				( tf.Module      ):







					"""simple docstring"""



					def __init__(self       ,			__lowercase		):
											super().__init__()
											__lowerCAmelCase							       =				tokenizer
											__lowerCAmelCase							       =				AutoConfig.from_pretrained(__lowercase		)
											__lowerCAmelCase							       =				TFGPTaLMHeadModel.from_config(__lowercase		)




					@tf.function(input_signature=(tf.TensorSpec((None,)       ,			tf.string       ,			name='''text'''		),)		)
					def 					_snake_case  (self       ,			__lowercase		):
											__lowerCAmelCase							       =				self.tokenizer(__lowercase		)
											__lowerCAmelCase							       =				tokenized['''input_ids'''].to_tensor()

											__lowerCAmelCase							       =				tf.cast(input_ids_dense > 0       ,			tf.intaa		)
											# input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

											__lowerCAmelCase							       =				self.model(input_ids=__lowercase       ,			attention_mask=__lowercase		)['''logits''']

											return outputs


@require_tf
@require_keras_nlp
class       a__				( unittest.TestCase      ):







				"""simple docstring"""



				def 					_snake_case  (self		):
										super().setUp()

										__lowerCAmelCase							       =				[GPTaTokenizer.from_pretrained(__lowercase		) for checkpoint in (TOKENIZER_CHECKPOINTS)]
										__lowerCAmelCase							       =				[TFGPTaTokenizer.from_pretrained(__lowercase		) for checkpoint in TOKENIZER_CHECKPOINTS]
										assert len(self.tokenizers		) == len(self.tf_tokenizers		)

										__lowerCAmelCase							       =				[
										    '''This is a straightforward English test sentence.''',
										    '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
										    '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
										    '''And some much more rare Chinese: 齉 堃 齉堃''',
										    '''Je vais aussi écrire en français pour tester les accents''',
										    '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
										]
										__lowerCAmelCase							       =				list(zip(self.test_sentences       ,			self.test_sentences[::-1]		)		)




				def 					_snake_case  (self		):
										for tokenizer, tf_tokenizer in zip(self.tokenizers       ,			self.tf_tokenizers		):
																for test_inputs in self.test_sentences:
																						__lowerCAmelCase							       =				tokenizer([test_inputs]       ,			return_tensors='''tf'''		)
																						__lowerCAmelCase							       =				tf_tokenizer([test_inputs]		)

																						for key in python_outputs.keys():
																												# convert them to numpy to avoid messing with ragged tensors
																												__lowerCAmelCase							       =				python_outputs[key].numpy()
																												__lowerCAmelCase							       =				tf_outputs[key].numpy()

																												self.assertTrue(tf.reduce_all(python_outputs_values.shape == tf_outputs_values.shape		)		)
																												self.assertTrue(tf.reduce_all(tf.cast(__lowercase       ,			tf.intaa		) == tf_outputs_values		)		)




				@slow
				def 					_snake_case  (self		):
										for tf_tokenizer in self.tf_tokenizers:
																__lowerCAmelCase							       =				tf.function(__lowercase		)
																for test_inputs in self.test_sentences:
																						__lowerCAmelCase							       =				tf.constant(__lowercase		)
																						__lowerCAmelCase							       =				compiled_tokenizer(__lowercase		)
																						__lowerCAmelCase							       =				tf_tokenizer(__lowercase		)

																						for key in eager_outputs.keys():
																												self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]		)		)




				@slow
				def 					_snake_case  (self		):
										for tf_tokenizer in self.tf_tokenizers:
																__lowerCAmelCase							       =				ModelToSave(tokenizer=__lowercase		)
																__lowerCAmelCase							       =				tf.convert_to_tensor([self.test_sentences[0]]		)
																__lowerCAmelCase							       =				model.serving(__lowercase		)  # Build model with some sample inputs
																with TemporaryDirectory() as tempdir:
																						__lowerCAmelCase							       =				Path(__lowercase		) / '''saved.model'''
																						tf.saved_model.save(__lowercase       ,			__lowercase       ,			signatures={'''serving_default''': model.serving}		)
																						__lowerCAmelCase							       =				tf.saved_model.load(__lowercase		)
																__lowerCAmelCase							       =				loaded_model.signatures['''serving_default'''](__lowercase		)['''output_0''']
																# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
																self.assertTrue(tf.reduce_all(out == loaded_output		)		)




				@slow
				def 					_snake_case  (self		):
										for tf_tokenizer in self.tf_tokenizers:
																__lowerCAmelCase							       =				tf.convert_to_tensor([self.test_sentences[0]]		)
																__lowerCAmelCase							       =				tf_tokenizer(__lowercase		)  # Build model with some sample inputs

																__lowerCAmelCase							       =				tf_tokenizer.get_config()
																__lowerCAmelCase							       =				TFGPTaTokenizer.from_config(__lowercase		)
																__lowerCAmelCase							       =				model_from_config(__lowercase		)

																for key in from_config_output.keys():
																						self.assertTrue(tf.reduce_all(from_config_output[key] == out[key]		)		)




				@slow
				def 					_snake_case  (self		):
										for tf_tokenizer in self.tf_tokenizers:
																# for the test to run
																__lowerCAmelCase							       =				12_31_23

																for max_length in [3, 5, 10_24]:
																						__lowerCAmelCase							       =				tf.convert_to_tensor([self.test_sentences[0]]		)
																						__lowerCAmelCase							       =				tf_tokenizer(__lowercase       ,			max_length=__lowercase		)

																						__lowerCAmelCase							       =				out['''input_ids'''].numpy().shape[1]

																						assert out_length == max_length

