





"""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      UpperCAmelCase__				(    ):







      '''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=5_12				,		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      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Optional[int]     ):







      '''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      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		Tuple     ):







      '''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      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		str				,		SCREAMING_SNAKE_CASE	:		Tuple				,		SCREAMING_SNAKE_CASE	:		Optional[Any]				,		SCREAMING_SNAKE_CASE	:		List[Any]				,		SCREAMING_SNAKE_CASE	:		Dict				,		SCREAMING_SNAKE_CASE	:		Union[str, Any]=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      UpperCAmelCase__				(    SCREAMING_SNAKE_CASE	:		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	:		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	:		Dict     ):
            # 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 json
from typing import Dict, List, Optional, Tuple, Union

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding, EncodedInput
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import PaddingStrategy, logging
from .tokenization_led import LEDTokenizer


SCREAMING_SNAKE_CASE__   =  logging.get_logger(__name__)


SCREAMING_SNAKE_CASE__   =  {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

SCREAMING_SNAKE_CASE__   =  {
    "vocab_file": {
        "allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json",
    },
    "merges_file": {
        "allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json",
    },
}

SCREAMING_SNAKE_CASE__   =  {
    "allenai/led-base-16384": 16_384,
}

class 						lowercase   (   _UpperCAmelCase   ):
      _SCREAMING_SNAKE_CASE							      = VOCAB_FILES_NAMES
      _SCREAMING_SNAKE_CASE							      = PRETRAINED_VOCAB_FILES_MAP
      _SCREAMING_SNAKE_CASE							      = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _SCREAMING_SNAKE_CASE							      = LEDTokenizer
      _SCREAMING_SNAKE_CASE							      = ['input_ids', 'attention_mask']





      def __init__(			self ,		lowercase=None ,		lowercase=None ,		lowercase=None ,		lowercase="replace" ,		lowercase="<s>" ,		lowercase="</s>" ,		lowercase="</s>" ,		lowercase="<s>" ,		lowercase="<unk>" ,		lowercase="<pad>" ,		lowercase="<mask>" ,		lowercase=False ,		lowercase=True ,		**lowercase ,		)						->	Any:
            super().__init__(
                lowercase ,		lowercase ,		tokenizer_file=lowercase ,		errors=lowercase ,		bos_token=lowercase ,		eos_token=lowercase ,		sep_token=lowercase ,		cls_token=lowercase ,		unk_token=lowercase ,		pad_token=lowercase ,		mask_token=lowercase ,		add_prefix_space=lowercase ,		trim_offsets=lowercase ,		**lowercase ,		)

            lowerCAmelCase    		=			json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()					)
            if pre_tok_state.get("""add_prefix_space""" ,		lowercase					) != add_prefix_space:
                  lowerCAmelCase    		=			getattr(lowercase ,		pre_tok_state.pop("""type"""					)					)
                  lowerCAmelCase    		=			add_prefix_space
                  lowerCAmelCase    		=			pre_tok_class(**lowercase					)

            lowerCAmelCase    		=			add_prefix_space

            # the pre_tokenizer is already updated in the GPT2TokenizerFast `__init__`
            lowerCAmelCase    		=			"""post_processor"""
            lowerCAmelCase    		=			getattr(self.backend_tokenizer ,		lowercase ,		lowercase					)
            if tokenizer_component_instance:
                  lowerCAmelCase    		=			json.loads(tokenizer_component_instance.__getstate__()					)

                  # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
                  if "sep" in state:
                        lowerCAmelCase    		=			tuple(state["""sep"""]					)
                  if "cls" in state:
                        lowerCAmelCase    		=			tuple(state["""cls"""]					)

                  lowerCAmelCase    		=			False

                  if state.get("""add_prefix_space""" ,		lowercase					) != add_prefix_space:
                        lowerCAmelCase    		=			add_prefix_space
                        lowerCAmelCase    		=			True

                  if state.get("""trim_offsets""" ,		lowercase					) != trim_offsets:
                        lowerCAmelCase    		=			trim_offsets
                        lowerCAmelCase    		=			True

                  if changes_to_apply:
                        lowerCAmelCase    		=			getattr(lowercase ,		state.pop("""type"""					)					)
                        lowerCAmelCase    		=			component_class(**lowercase					)
                        setattr(self.backend_tokenizer ,		lowercase ,		lowercase					)





      @property
      # Copied from transformers.models.bart.tokenization_bart_fast.BartTokenizerFast.mask_token with BART->LED
      def        _snake_case     (			self					)						->	str:
            if self._mask_token is None:
                  if self.verbose:
                        logger.error("""Using mask_token, but it is not set yet."""					)
                  return None
            return str(self._mask_token					)





      @mask_token.setter
      def        _snake_case     (			self ,		lowercase					)						->	Optional[int]:
            lowerCAmelCase    		=			AddedToken(lowercase ,		lstrip=lowercase ,		rstrip=lowercase					) if isinstance(lowercase ,		lowercase					) else value
            lowerCAmelCase    		=			value





      def        _snake_case     (			self ,		*lowercase ,		**lowercase					)						->	BatchEncoding:
            lowerCAmelCase    		=			kwargs.get("""is_split_into_words""" ,		lowercase					)

            if is_split_into_words and not self.add_prefix_space:
                  raise ValueError(
                      f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
                      """to use it with pretokenized inputs."""					)

            return super()._batch_encode_plus(*lowercase ,		**lowercase					)





      def        _snake_case     (			self ,		*lowercase ,		**lowercase					)						->	BatchEncoding:
            lowerCAmelCase    		=			kwargs.get("""is_split_into_words""" ,		lowercase					)

            if is_split_into_words and not self.add_prefix_space:
                  raise ValueError(
                      f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
                      """to use it with pretokenized inputs."""					)

            return super()._encode_plus(*lowercase ,		**lowercase					)





      def        _snake_case     (			self ,		lowercase ,		lowercase = None					)						->	Tuple[str]:
            lowerCAmelCase    		=			self._tokenizer.model.save(lowercase ,		name=lowercase					)
            return tuple(lowercase					)





      def        _snake_case     (			self ,		lowercase ,		lowercase=None					)						->	Any:
            lowerCAmelCase    		=			[self.bos_token_id] + token_ids_a + [self.eos_token_id]
            if token_ids_a is None:
                  return output

            return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]





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

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





      def        _snake_case     (			self ,		lowercase ,		lowercase = None ,		lowercase = PaddingStrategy.DO_NOT_PAD ,		lowercase = None ,		lowercase = None ,		)						->	dict:
            lowerCAmelCase    		=			super()._pad(
                encoded_inputs=lowercase ,		max_length=lowercase ,		padding_strategy=lowercase ,		pad_to_multiple_of=lowercase ,		return_attention_mask=lowercase ,		)

            # Load from model defaults
            if return_attention_mask is None:
                  lowerCAmelCase    		=			"""attention_mask""" in self.model_input_names

            if return_attention_mask and "global_attention_mask" in encoded_inputs:
                  lowerCAmelCase    		=			encoded_inputs[self.model_input_names[0]]
                  # `global_attention_mask` need to have the same length as other (sequential) inputs.
                  lowerCAmelCase    		=			len(encoded_inputs["""global_attention_mask"""]					) != len(lowercase					)

                  if needs_to_be_padded:
                        lowerCAmelCase    		=			len(lowercase					) - len(encoded_inputs["""global_attention_mask"""]					)

                        if self.padding_side == "right":
                              # Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
                              lowerCAmelCase    		=			(
                                  encoded_inputs["""global_attention_mask"""] + [-1] * difference
                              )
                        elif self.padding_side == "left":
                              lowerCAmelCase    		=			[-1] * difference + encoded_inputs[
                                  """global_attention_mask"""
                              ]
                        else:
                              raise ValueError("""Invalid padding strategy:""" + str(self.padding_side					)					)

            return encoded_inputs






