





'''simple docstring'''


import argparse
import logging
import os
import re

import tensorflow as tf

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


_SCREAMING_SNAKE_CASE	 =							logging.getLogger(__name__)

_SCREAMING_SNAKE_CASE	 =							tf.data.AUTOTUNE


def 	__a():



    '''simple docstring'''





    _lowerCAmelCase				=      argparse.ArgumentParser(description="Train a masked language model on TPU."    )
    parser.add_argument(
        "--pretrained_model_config"		,       type=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       help="Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."		,       )

    parser.add_argument(
        "--gcp_project"		,       type=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       default=2**18		,       help="Size of the shuffle buffer (in samples)"		,       )

    parser.add_argument(
        "--eval_dataset"		,       type=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       default=1		,       help="Number of epochs to train for."		,       )

    parser.add_argument(
        "--learning_rate"		,       type=SCREAMING_SNAKE_CASE_		,       default=1e-4		,       help="Learning rate to use for training."		,       )

    parser.add_argument(
        "--weight_decay_rate"		,       type=SCREAMING_SNAKE_CASE_		,       default=1e-3		,       help="Weight decay rate to use for training."		,       )

    parser.add_argument(
        "--max_length"		,       type=SCREAMING_SNAKE_CASE_		,       default=512		,       help="Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py"		,       )

    parser.add_argument(
        "--mlm_probability"		,       type=SCREAMING_SNAKE_CASE_		,       default=0.15		,       help="Fraction of tokens to mask during training."		,       )

    parser.add_argument("--output_dir"		,       type=SCREAMING_SNAKE_CASE_		,       required=SCREAMING_SNAKE_CASE_		,       help="Path to save model checkpoints to."    )
    parser.add_argument("--hub_model_id"		,       type=SCREAMING_SNAKE_CASE_		,       help="Model ID to upload to on the Hugging Face Hub."    )

    _lowerCAmelCase				=      parser.parse_args()
    return args


def 	__a(SCREAMING_SNAKE_CASE_  :					List[Any]    ):



    '''simple docstring'''





    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(SCREAMING_SNAKE_CASE_    )
    tf.tpu.experimental.initialize_tpu_system(SCREAMING_SNAKE_CASE_    )

    return tpu


def 	__a(SCREAMING_SNAKE_CASE_  :					int    ):



    '''simple docstring'''





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

    return num_samples


def 	__a(SCREAMING_SNAKE_CASE_  :					Tuple		,       SCREAMING_SNAKE_CASE_  :					str		,       SCREAMING_SNAKE_CASE_  :					Any		,       SCREAMING_SNAKE_CASE_  :					List[str]		,       SCREAMING_SNAKE_CASE_  :					List[Any]		,       SCREAMING_SNAKE_CASE_  :					Tuple=None    ):



    '''simple docstring'''





    _lowerCAmelCase				=      count_samples(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      tf.data.Dataset.from_tensor_slices(SCREAMING_SNAKE_CASE_    )
    if shuffle:
        _lowerCAmelCase				=      dataset.shuffle(len(SCREAMING_SNAKE_CASE_    )    )
    _lowerCAmelCase				=      tf.data.TFRecordDataset(SCREAMING_SNAKE_CASE_		,       num_parallel_reads=SCREAMING_SNAKE_CASE_    )
    # 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(SCREAMING_SNAKE_CASE_    )    )
    _lowerCAmelCase				=      dataset.map(SCREAMING_SNAKE_CASE_		,       num_parallel_calls=SCREAMING_SNAKE_CASE_    )
    if shuffle:
        assert shuffle_buffer_size is not None
        _lowerCAmelCase				=      dataset.shuffle(args.shuffle_buffer_size    )
    _lowerCAmelCase				=      dataset.batch(SCREAMING_SNAKE_CASE_		,       drop_remainder=SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      dataset.map(SCREAMING_SNAKE_CASE_		,       num_parallel_calls=SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      dataset.prefetch(SCREAMING_SNAKE_CASE_    )
    return dataset


def 	__a(SCREAMING_SNAKE_CASE_  :					Optional[Any]    ):



    '''simple docstring'''





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

    _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(SCREAMING_SNAKE_CASE_    )
        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=SCREAMING_SNAKE_CASE_		,       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=SCREAMING_SNAKE_CASE_		,       metrics=["accuracy"]    )

    def decode_fn(SCREAMING_SNAKE_CASE_  :					Optional[int]    ):
        _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(SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    )

    # 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=SCREAMING_SNAKE_CASE_		,       mlm_probability=args.mlm_probability		,       mlm=SCREAMING_SNAKE_CASE_		,       return_tensors="tf"    )

    def mask_with_collator(SCREAMING_SNAKE_CASE_  :					str    ):
        # 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(SCREAMING_SNAKE_CASE_    )		,       mask_token_id=tokenizer.mask_token_id		,       special_tokens_mask=SCREAMING_SNAKE_CASE_		,       )
        return batch

    _lowerCAmelCase				=      args.per_replica_batch_size * strategy.num_replicas_in_sync

    _lowerCAmelCase				=      prepare_dataset(
        SCREAMING_SNAKE_CASE_		,       decode_fn=SCREAMING_SNAKE_CASE_		,       mask_fn=SCREAMING_SNAKE_CASE_		,       batch_size=SCREAMING_SNAKE_CASE_		,       shuffle=SCREAMING_SNAKE_CASE_		,       shuffle_buffer_size=args.shuffle_buffer_size		,       )

    _lowerCAmelCase				=      prepare_dataset(
        SCREAMING_SNAKE_CASE_		,       decode_fn=SCREAMING_SNAKE_CASE_		,       mask_fn=SCREAMING_SNAKE_CASE_		,       batch_size=SCREAMING_SNAKE_CASE_		,       shuffle=SCREAMING_SNAKE_CASE_		,       )

    _lowerCAmelCase				=      []
    if args.hub_model_id:
        callbacks.append(
            PushToHubCallback(output_dir=args.output_dir		,       hub_model_id=args.hub_model_id		,       tokenizer=SCREAMING_SNAKE_CASE_    )    )

    model.fit(
        SCREAMING_SNAKE_CASE_		,       validation_data=SCREAMING_SNAKE_CASE_		,       epochs=args.num_epochs		,       callbacks=SCREAMING_SNAKE_CASE_		,       )

    model.save_pretrained(args.output_dir    )


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE	 =							parse_args()
       main(args)









'''simple docstring'''


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


_SCREAMING_SNAKE_CASE	 =							logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE	 =							"▁"

_SCREAMING_SNAKE_CASE	 =							{"vocab_file": "sentencepiece.bpe.model"}

_SCREAMING_SNAKE_CASE	 =							{
    "vocab_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model",
    }
}

_SCREAMING_SNAKE_CASE	 =							{
    "facebook/xglm-564M": 20_48,
}

class    lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase						:							int         =					VOCAB_FILES_NAMES
 __lowerCamelCase						:							Optional[int]         =					PRETRAINED_VOCAB_FILES_MAP
 __lowerCamelCase						:							Tuple         =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCamelCase						:							str         =					["input_ids", "attention_mask"]




 def __init__(	self			,     _lowerCAmelCase			,     _lowerCAmelCase="<s>"			,     _lowerCAmelCase="</s>"			,     _lowerCAmelCase="</s>"			,     _lowerCAmelCase="<s>"			,     _lowerCAmelCase="<unk>"			,     _lowerCAmelCase="<pad>"			,     _lowerCAmelCase = None			,     **_lowerCAmelCase			,     )    ->							None:
     _lowerCAmelCase				=      {} if sp_model_kwargs is None else sp_model_kwargs

     # Compatibility with the original tokenizer
     _lowerCAmelCase				=      7
     _lowerCAmelCase				=      [f'''<madeupword{i}>''' for i in range(self.num_madeup_words		)]

     _lowerCAmelCase				=      kwargs.get("additional_special_tokens"			,     []		)
     kwargs["additional_special_tokens"] += [
         word for word in madeup_words if word not in kwargs["additional_special_tokens"]
     ]

     super().__init__(
         bos_token=_lowerCAmelCase			,     eos_token=_lowerCAmelCase			,     unk_token=_lowerCAmelCase			,     sep_token=_lowerCAmelCase			,     cls_token=_lowerCAmelCase			,     pad_token=_lowerCAmelCase			,     sp_model_kwargs=self.sp_model_kwargs			,     **_lowerCAmelCase			,     )

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

     # Original fairseq vocab and spm vocab must be "aligned":
     # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
     # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
     # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
     # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

     # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
     _lowerCAmelCase				=      1

     # Mimic fairseq token-to-id alignment for the first 4 token
     _lowerCAmelCase				=      {"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

     _lowerCAmelCase				=      len(self.sp_model		)
     _lowerCAmelCase				=      {f'''<madeupword{i}>''': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words		)}
     self.fairseq_tokens_to_ids.update(_lowerCAmelCase		)

     _lowerCAmelCase				=      {v: k for k, v in self.fairseq_tokens_to_ids.items()}




 def __getstate__(	self		)    ->							List[str]:
     _lowerCAmelCase				=      self.__dict__.copy()
     _lowerCAmelCase				=      None
     _lowerCAmelCase				=      self.sp_model.serialized_model_proto()
     return state




 def __setstate__(	self			,     _lowerCAmelCase		)    ->							Optional[int]:
     _lowerCAmelCase				=      d

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

     _lowerCAmelCase				=      spm.SentencePieceProcessor(**self.sp_model_kwargs		)
     self.sp_model.LoadFromSerializedProto(self.sp_model_proto		)




 def 							_snake_case   (	self			,     _lowerCAmelCase			,     _lowerCAmelCase = None		)    ->							List[int]:

     if token_ids_a is None:
         return [self.sep_token_id] + token_ids_a
     _lowerCAmelCase				=      [self.sep_token_id]
     return sep + token_ids_a + sep + sep + token_ids_a




 def 							_snake_case   (	self			,     _lowerCAmelCase			,     _lowerCAmelCase = None			,     _lowerCAmelCase = False		)    ->							List[int]:

     if already_has_special_tokens:
         return super().get_special_tokens_mask(
             token_ids_a=_lowerCAmelCase			,     token_ids_a=_lowerCAmelCase			,     already_has_special_tokens=_lowerCAmelCase		)

     if token_ids_a is None:
         return [1] + ([0] * len(_lowerCAmelCase		))
     return [1] + ([0] * len(_lowerCAmelCase		)) + [1, 1] + ([0] * len(_lowerCAmelCase		))




 def 							_snake_case   (	self			,     _lowerCAmelCase			,     _lowerCAmelCase = None		)    ->							List[int]:
     _lowerCAmelCase				=      [self.sep_token_id]

     if token_ids_a is None:
         return len(sep + token_ids_a		) * [0]
     return len(sep + token_ids_a + sep + sep + token_ids_a		) * [0]




 @property
 def 							_snake_case   (	self		)    ->							str:
     return len(self.sp_model		) + self.fairseq_offset + self.num_madeup_words




 def 							_snake_case   (	self		)    ->							Any:
     _lowerCAmelCase				=      {self.convert_ids_to_tokens(_lowerCAmelCase		): i for i in range(self.vocab_size		)}
     vocab.update(self.added_tokens_encoder		)
     return vocab




 def 							_snake_case   (	self			,     _lowerCAmelCase		)    ->							List[str]:
     return self.sp_model.encode(_lowerCAmelCase			,     out_type=_lowerCAmelCase		)




 def 							_snake_case   (	self			,     _lowerCAmelCase		)    ->							List[str]:
     if token in self.fairseq_tokens_to_ids:
         return self.fairseq_tokens_to_ids[token]
     _lowerCAmelCase				=      self.sp_model.PieceToId(_lowerCAmelCase		)

     # Need to return unknown token if the SP model returned 0
     return spm_id + self.fairseq_offset if spm_id else self.unk_token_id




 def 							_snake_case   (	self			,     _lowerCAmelCase		)    ->							List[Any]:
     if index in self.fairseq_ids_to_tokens:
         return self.fairseq_ids_to_tokens[index]
     return self.sp_model.IdToPiece(index - self.fairseq_offset		)




 def 							_snake_case   (	self			,     _lowerCAmelCase		)    ->							Union[str, Any]:
     _lowerCAmelCase				=      "".join(_lowerCAmelCase		).replace(_lowerCAmelCase			,     " "		).strip()
     return out_string





 def 							_snake_case   (	self			,     _lowerCAmelCase			,     _lowerCAmelCase = None		)    ->							Tuple[str]:
     if not os.path.isdir(_lowerCAmelCase		):
         logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''		)
         return
     _lowerCAmelCase				=      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:
             _lowerCAmelCase				=      self.sp_model.serialized_model_proto()
             fi.write(_lowerCAmelCase		)

     return (out_vocab_file,)


