




import argparse
import logging
import os
import re

import tensorflow as tf

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


__lowerCAmelCase	:       Dict     =	logging.getLogger(__name__)

__lowerCAmelCase	:       int     =	tf.data.AUTOTUNE




def     __magic_name__       (				):






			'''simple docstring'''







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

			a					=    parser.parse_args()
			return args




def     __magic_name__       (				A	:      Tuple     ):






			'''simple docstring'''







			try:
						if args.tpu_name:
									a					=    tf.distribute.cluster_resolver.TPUClusterResolver(
									    args.tpu_name,						zone=args.tpu_zone,						project=args.gcp_project     )
						else:
									a					=    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     __magic_name__       (				A	:      int     ):






			'''simple docstring'''







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

			return num_samples




def     __magic_name__       (				A	:      Tuple,						A	:      List[str],						A	:      str,						A	:      Optional[int],						A	:      List[str],						A	:      Optional[int]=None     ):






			'''simple docstring'''







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




def     __magic_name__       (				A	:      List[str]     ):






			'''simple docstring'''







			if not args.no_tpu:
						a					=    initialize_tpu(UpperCAmelCase__     )
						a					=    tf.distribute.TPUStrategy(UpperCAmelCase__     )
			else:
						a					=    tf.distribute.OneDeviceStrategy(device="/gpu:0"     )

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

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

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

			a					=    count_samples(UpperCAmelCase__     )

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

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

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

			a					=    args.per_replica_batch_size * strategy.num_replicas_in_sync

			a					=    prepare_dataset(
			    UpperCAmelCase__,						decode_fn=UpperCAmelCase__,						mask_fn=UpperCAmelCase__,						batch_size=UpperCAmelCase__,						shuffle=UpperCAmelCase__,						shuffle_buffer_size=args.shuffle_buffer_size,						)

			a					=    prepare_dataset(
			    UpperCAmelCase__,						decode_fn=UpperCAmelCase__,						mask_fn=UpperCAmelCase__,						batch_size=UpperCAmelCase__,						shuffle=UpperCAmelCase__,						)

			a					=    []
			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	:       List[Any]     =	parse_args()
				main(args)












import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__lowerCAmelCase	:       int     =	logging.get_logger(__name__)


__lowerCAmelCase	:       Optional[int]     =	'▁'

__lowerCAmelCase	:       Union[str, Any]     =	{'vocab_file': 'spiece.model'}

__lowerCAmelCase	:       int     =	{
    'vocab_file': {
        'google/reformer-crime-and-punishment': (
            'https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/spiece.model'
        )
    }
}

__lowerCAmelCase	:       Any     =	{
    'google/reformer-crime-and-punishment': 52_4288,
}




class 					snake_case__							(_UpperCamelCase     ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_ :							Dict											=      VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE_ :							int											=      PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE_ :							List[Any]											=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE_ :							Optional[int]											=      ["""input_ids""", """attention_mask"""]





				def __init__(   self       :  List[Any]      ,   __lowerCamelCase       :  List[str]      ,   __lowerCamelCase       :  Dict="</s>"      ,   __lowerCamelCase       :  Dict="<unk>"      ,   __lowerCamelCase       :  Dict=[]      ,   __lowerCamelCase       :  Optional[Dict[str, Any]] = None      ,   **__lowerCamelCase       :  Dict      ,   )      ->  None:
							a					=    {} if sp_model_kwargs is None else sp_model_kwargs

							super().__init__(
							    eos_token=__lowerCamelCase      ,   unk_token=__lowerCamelCase      ,   additional_special_tokens=__lowerCamelCase      ,   sp_model_kwargs=self.sp_model_kwargs      ,   **__lowerCamelCase      ,   )

							a					=    vocab_file
							a					=    spm.SentencePieceProcessor(**self.sp_model_kwargs							)
							self.sp_model.Load(__lowerCamelCase							)





				@property
				def 							__UpperCAmelCase							(   self       :  Optional[int]							)      ->  int:
							return self.sp_model.get_piece_size()





				def 							__UpperCAmelCase							(   self       :  Tuple							)      ->  Dict[str, int]:
							a					=    {self.convert_ids_to_tokens(__lowerCamelCase							): i for i in range(self.vocab_size							)}
							vocab.update(self.added_tokens_encoder							)
							return vocab





				def __getstate__(   self       :  Optional[Any]							)      ->  Optional[Any]:
							a					=    self.__dict__.copy()
							a					=    None
							return state





				def __setstate__(   self       :  str      ,   __lowerCamelCase       :  Tuple							)      ->  List[Any]:
							a					=    d

							# for backward compatibility
							if not hasattr(self      ,   "sp_model_kwargs"							):
										a					=    {}

							a					=    spm.SentencePieceProcessor(**self.sp_model_kwargs							)
							self.sp_model.Load(self.vocab_file							)





				def 							__UpperCAmelCase							(   self       :  int      ,   __lowerCamelCase       :  str							)      ->  List[str]:
							return self.sp_model.encode(__lowerCamelCase      ,   out_type=__lowerCamelCase							)





				def 							__UpperCAmelCase							(   self       :  List[Any]      ,   __lowerCamelCase       :  Dict							)      ->  Any:
							return self.sp_model.piece_to_id(__lowerCamelCase							)





				def 							__UpperCAmelCase							(   self       :  int      ,   __lowerCamelCase       :  Union[str, Any]							)      ->  str:
							if index < self.sp_model.get_piece_size():
										a					=    self.sp_model.IdToPiece(__lowerCamelCase							)
							return token





				def 							__UpperCAmelCase							(   self       :  Optional[int]      ,   __lowerCamelCase       :  Optional[Any]							)      ->  List[Any]:
							a					=    []
							a					=    ""
							for token in tokens:
										# make sure that special tokens are not decoded using sentencepiece model
										if token in self.all_special_tokens:
													out_string += self.sp_model.decode(__lowerCamelCase							) + token
													a					=    []
										else:
													current_sub_tokens.append(__lowerCamelCase							)
							out_string += self.sp_model.decode(__lowerCamelCase							)
							return out_string.strip()





				def 							__UpperCAmelCase							(   self       :  Any      ,   __lowerCamelCase       :  str      ,   __lowerCamelCase       :  Optional[str] = None							)      ->  Tuple[str]:
							if not os.path.isdir(__lowerCamelCase							):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""							)
										return
							a					=    os.path.join(
							    __lowerCamelCase      ,   (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]							)

							if os.path.abspath(self.vocab_file							) != os.path.abspath(__lowerCamelCase							) and os.path.isfile(self.vocab_file							):
										copyfile(self.vocab_file      ,   __lowerCamelCase							)
							elif not os.path.isfile(self.vocab_file							):
										with open(__lowerCamelCase      ,   "wb"							) as fi:
													a					=    self.sp_model.serialized_model_proto()
													fi.write(__lowerCamelCase							)

							return (out_vocab_file,)






