"""simple docstring"""




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

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
      from .tokenization_big_bird import BigBirdTokenizer
else:
      UpperCamelCase__    							= None

UpperCamelCase__    							= logging.get_logger(__name__)
UpperCamelCase__    							= {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

UpperCamelCase__    							= {
    """vocab_file""": {
        """google/bigbird-roberta-base""": """https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model""",
        """google/bigbird-roberta-large""": (
            """https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model"""
        ),
        """google/bigbird-base-trivia-itc""": (
            """https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model"""
        ),
    },
    """tokenizer_file""": {
        """google/bigbird-roberta-base""": (
            """https://huggingface.co/google/bigbird-roberta-base/resolve/main/tokenizer.json"""
        ),
        """google/bigbird-roberta-large""": (
            """https://huggingface.co/google/bigbird-roberta-large/resolve/main/tokenizer.json"""
        ),
        """google/bigbird-base-trivia-itc""": (
            """https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/tokenizer.json"""
        ),
    },
}

UpperCamelCase__    							= {
    """google/bigbird-roberta-base""": 4_0_9_6,
    """google/bigbird-roberta-large""": 4_0_9_6,
    """google/bigbird-base-trivia-itc""": 4_0_9_6,
}


UpperCamelCase__    							= """▁"""





class 					A							(    SCREAMING_SNAKE_CASE_					):
   __UpperCAmelCase			:			List[str]              =      VOCAB_FILES_NAMES
   __UpperCAmelCase			:			int              =      PRETRAINED_VOCAB_FILES_MAP
   __UpperCAmelCase			:			List[str]              =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   __UpperCAmelCase			:			Optional[Any]              =      BigBirdTokenizer
   __UpperCAmelCase			:			Optional[Any]              =      ['input_ids', 'attention_mask']
   __UpperCAmelCase			:			Optional[int]              =      []


   def __init__(self					:						int	,  __UpperCAmelCase					:						Dict=None	,  __UpperCAmelCase					:						int=None	,  __UpperCAmelCase					:						List[Any]="<unk>"	,  __UpperCAmelCase					:						Optional[Any]="<s>"	,  __UpperCAmelCase					:						List[Any]="</s>"	,  __UpperCAmelCase					:						Optional[Any]="<pad>"	,  __UpperCAmelCase					:						Union[str, Any]="[SEP]"	,  __UpperCAmelCase					:						Optional[int]="[MASK]"	,  __UpperCAmelCase					:						Dict="[CLS]"	,  **__UpperCAmelCase					:						str	,  )						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else bos_token
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else eos_token
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else unk_token
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else pad_token
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else cls_token
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else sep_token

         # Mask token behave like a normal word, i.e. include the space before it
         UpperCAmelCase__								=  AddedToken(snake_case__	,  lstrip=snake_case__	,  rstrip=snake_case__		) if isinstance(snake_case__	,  snake_case__		) else mask_token

         super().__init__(
             snake_case__	,  tokenizer_file=snake_case__	,  bos_token=snake_case__	,  eos_token=snake_case__	,  unk_token=snake_case__	,  sep_token=snake_case__	,  pad_token=snake_case__	,  cls_token=snake_case__	,  mask_token=snake_case__	,  **snake_case__	,  )

         UpperCAmelCase__								=  vocab_file
         UpperCAmelCase__								=  False if not self.vocab_file else True


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any] = None		)						->    Any:
         """simple docstring"""



         UpperCAmelCase__								=  [self.sep_token_id]
         UpperCAmelCase__								=  [self.cls_token_id]
         if token_ids_a is None:
               return cls + token_ids_a + sep
         return cls + token_ids_a + sep + token_ids_a + sep


   def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[Any] = None	,  __UpperCAmelCase					:						Any = False		)						->    List[Any]:
         """simple docstring"""



         if already_has_special_tokens:
               if token_ids_a is not None:
                     raise ValueError(
                         "You should not supply a second sequence if the provided sequence of "
                         "ids is already formatted with special tokens for the model."		)
               return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

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


   def       lowercase_				(self					:						Union[str, Any]	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						List[str] = None		)						->    Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  [self.sep_token_id]
         UpperCAmelCase__								=  [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		) * [0] + len(token_ids_a + sep		) * [1]





   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Any = None		)						->    Tuple:
         """simple docstring"""



         if not self.can_save_slow_tokenizer:
               raise ValueError(
                   "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                   "tokenizer."		)

         if not os.path.isdir(snake_case__		):
               logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""		)
               return
         UpperCAmelCase__								=  os.path.join(
             snake_case__	,  (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]		)

         if os.path.abspath(self.vocab_file		) != os.path.abspath(snake_case__		):
               copyfile(self.vocab_file	,  snake_case__		)

         return (out_vocab_file,)




import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


UpperCamelCase__    							= logging.getLogger(__name__)




def        lowerCAmelCase_			(					)					->    Union[str, Any]:







      '''simple docstring'''


      UpperCAmelCase__								=  argparse.ArgumentParser(
          description="Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."							)
      parser.add_argument(
          "--dataset_name",  type=__A,  default="wikitext",  help="Name of the training. Explore datasets at: hf.co/datasets.",  )
      parser.add_argument(
          "--dataset_config",  type=__A,  default="wikitext-103-raw-v1",  help="Configuration name of the dataset."							)
      parser.add_argument(
          "--tokenizer_name_or_path",  type=__A,  default="sayakpaul/unigram-tokenizer-wikitext",  help="Tokenizer identifier. Can be a local filepath or a Hub identifier.",  )
      parser.add_argument(
          "--shard_size",  type=__A,  default=1_000,  help="Number of entries to go in a single shard.",  )
      parser.add_argument("--split",  type=__A,  default="train",  choices=["train", "test", "validation"]							)
      parser.add_argument(
          "--limit",  default=__A,  type=__A,  help="Limit the number of shards (used for debugging).",  )
      parser.add_argument(
          "--max_length",  type=__A,  default=512,  help="Maximum sequence length. For training on TPUs, it helps to have a maximum"
          " sequence length that is a multiple of 8.",  )
      parser.add_argument(
          "--output_dir",  default="tf-tpu",  type=__A,  help="Output directory where the TFRecord shards will be saved. If the"
          " path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"
          " shards will be directly saved to a Google Cloud Storage bucket.",  )

      UpperCAmelCase__								=  parser.parse_args()
      return args




def        lowerCAmelCase_			(					__A							)					->    Optional[int]:







      '''simple docstring'''


      def fn(__A							):
            return tokenizer(examples["text"]							)

      return fn




def        lowerCAmelCase_			(					__A							)					->    List[Any]:







      '''simple docstring'''


      UpperCAmelCase__								=  []
      for i in range(len(tokenized_data["input_ids"]							)							):
            UpperCAmelCase__								=  {
                "input_ids": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["input_ids"][i]							)							),
                "attention_mask": tf.train.Feature(
                    intaa_list=tf.train.IntaaList(value=tokenized_data["attention_mask"][i]							)							),
            }
            UpperCAmelCase__								=  tf.train.Features(feature=__A							)
            UpperCAmelCase__								=  tf.train.Example(features=__A							)
            UpperCAmelCase__								=  example.SerializeToString()
            records.append(__A							)
      return records




def        lowerCAmelCase_			(					__A							)					->    Union[str, Any]:







      '''simple docstring'''


      UpperCAmelCase__								=  datasets.load_dataset(args.dataset_name,  args.dataset_config,  split=args.split							)

      if args.limit is not None:
            UpperCAmelCase__								=  min(len(__A							),  args.limit							)
            UpperCAmelCase__								=  dataset.select(range(__A							)							)
            print(f"""Limiting the dataset to {args.limit} entries."""							)

      UpperCAmelCase__								=  AutoTokenizer.from_pretrained(args.tokenizer_name_or_path							)

      # Handle output directory creation.
      # For serializing into a Google Cloud Storage Bucket, one needs to first
      # create a bucket.
      if "gs" not in args.output_dir:
            if not os.path.exists(args.output_dir							):
                  os.makedirs(args.output_dir							)
            UpperCAmelCase__								=  os.path.join(args.output_dir,  args.split							)
            if not os.path.exists(__A							):
                  os.makedirs(__A							)
      else:
            UpperCAmelCase__								=  os.path.join(args.output_dir,  args.split							)

      # Tokenize the whole dataset at once.
      UpperCAmelCase__								=  tokenize_function(__A							)
      UpperCAmelCase__								=  dataset.map(__A,  batched=__A,  num_proc=4,  remove_columns=["text"]							)

      # We need to concatenate all our texts together, and then split the result
      # into chunks of a fixed size, which we will call block_size. To do this, we
      # will use the map method again, with the option batched=True. When we use batched=True,
      # the function we pass to map() will be passed multiple inputs at once, allowing us
      # to group them into more or fewer examples than we had in the input.
      # This allows us to create our new fixed-length samples. The advantage of this
      # method is that we don't lose a whole lot of content from the dataset compared to the
      # case where we simply tokenize with a pre-defined max_length.

      def group_texts(__A							):
            # Concatenate all texts.
            UpperCAmelCase__								=  {k: sum(examples[k],  []							) for k in examples.keys()}
            UpperCAmelCase__								=  len(concatenated_examples[list(examples.keys()							)[0]]							)
            # We drop the small remainder, though you could add padding instead if the model supports it
            # In this, as in all things, we advise you to follow your heart 🫀
            UpperCAmelCase__								=  (total_length // args.max_length) * args.max_length
            # Split by chunks of max_len.
            UpperCAmelCase__								=  {
                k: [t[i : i + args.max_length] for i in range(0,  __A,  args.max_length							)]
                for k, t in concatenated_examples.items()
            }
            return result

      UpperCAmelCase__								=  dataset_tokenized.map(__A,  batched=__A,  batch_size=1_000,  num_proc=4							)

      UpperCAmelCase__								=  0
      UpperCAmelCase__								=  0
      for shard in range(0,  len(__A							),  args.shard_size							):
            UpperCAmelCase__								=  grouped_dataset[shard : shard + args.shard_size]
            UpperCAmelCase__								=  len(dataset_snapshot["input_ids"]							)
            UpperCAmelCase__								=  os.path.join(__A,  f"""dataset-{shard_count}-{records_containing}.tfrecord"""							)
            UpperCAmelCase__								=  get_serialized_examples(__A							)

            with tf.io.TFRecordWriter(__A							) as out_file:
                  for i in range(len(__A							)							):
                        UpperCAmelCase__								=  serialized_examples[i]
                        out_file.write(__A							)
                  print("Wrote file {} containing {} records".format(__A,  __A							)							)

            shard_count += 1
            total_records += records_containing

      with open(f"""split-{args.split}-records-count.txt""",  "w"							) as f:
            print(f"""Total {args.split} records: {total_records}""",  file=__A							)


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



