



"""simple docstring"""







import json
import os
from typing import Dict, List, Optional, Tuple

import regex as re

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


lowercase__		:				Optional[Any]          =	logging.get_logger(__name__)


lowercase__		:				List[Any]          =	{
    """vocab_file""": """vocab.json""",
    """merges_file""": """merges.txt""",
    """tokenizer_config_file""": """tokenizer_config.json""",
}

lowercase__		:				Tuple          =	{
    """vocab_file""": {
        """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json"""
    },
    """merges_file""": {
        """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt"""
    },
    """tokenizer_config_file""": {
        """facebook/blenderbot_small-90M""": (
            """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json"""
        )
    },
}

lowercase__		:				List[str]          =	{"""facebook/blenderbot_small-90M""": 5_1_2}

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple						)						->   Any:

  """simple docstring"""




  lowerCAmelCase_   :				str								= set()
  lowerCAmelCase_   :				List[Any]								= word[0]
  for char in word[1:]:
    pairs.add((prev_char, char)						)
    lowerCAmelCase_   :				int								= char

  lowerCAmelCase_   :				List[Any]								= set(lowerCAmelCase__						)
  return pairs







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     VOCAB_FILES_NAMES
      _SCREAMING_SNAKE_CASE  				=     PRETRAINED_VOCAB_FILES_MAP
      _SCREAMING_SNAKE_CASE  				=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _SCREAMING_SNAKE_CASE  				=     ["""input_ids""", """attention_mask"""]



      def __init__(      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      int="__start__"    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]="__end__"    ,			SCREAMING_SNAKE_CASE_				:      str="__unk__"    ,			SCREAMING_SNAKE_CASE_				:      int="__null__"    ,			**SCREAMING_SNAKE_CASE_				:      Any    ,			):
        super().__init__(unk_token=SCREAMING_SNAKE_CASE_    ,			bos_token=SCREAMING_SNAKE_CASE_    ,			eos_token=SCREAMING_SNAKE_CASE_    ,			pad_token=SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_       )

        with open(SCREAMING_SNAKE_CASE_    ,			encoding='utf-8'       ) as vocab_handle:
          lowerCAmelCase_   :				Tuple								= json.load(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[str]								= {v: k for k, v in self.encoder.items()}
        with open(SCREAMING_SNAKE_CASE_    ,			encoding='utf-8'       ) as merges_handle:
          lowerCAmelCase_   :				Tuple								= merges_handle.read().split('\n'       )[1:-1]
        lowerCAmelCase_   :				Optional[Any]								= [tuple(merge.split()       ) for merge in merges]
        lowerCAmelCase_   :				Any								= dict(zip(SCREAMING_SNAKE_CASE_    ,			range(len(SCREAMING_SNAKE_CASE_       )       )       )       )
        lowerCAmelCase_   :				Any								= {}



      @property
      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]       ):
        return len(self.encoder       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        return dict(self.encoder    ,			**self.added_tokens_encoder       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        if token in self.cache:
          return self.cache[token]
        lowerCAmelCase_   :				Union[str, Any]								= re.sub('([.,!?()])'    ,			r' \1'    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Dict								= re.sub('(\')'    ,			r' \1 '    ,			SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				int								= re.sub(r'\s{2,}'    ,			' '    ,			SCREAMING_SNAKE_CASE_       )
        if "\n" in token:
          lowerCAmelCase_   :				Tuple								= token.replace('\n'    ,			' __newln__'       )

        lowerCAmelCase_   :				Any								= token.split(' '       )
        lowerCAmelCase_   :				str								= []
        for token in tokens:
          if not len(SCREAMING_SNAKE_CASE_       ):
            continue

          lowerCAmelCase_   :				Dict								= token.lower()
          lowerCAmelCase_   :				int								= tuple(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Union[str, Any]								= tuple(list(word[:-1]       ) + [word[-1] + '</w>']       )
          lowerCAmelCase_   :				Union[str, Any]								= get_pairs(SCREAMING_SNAKE_CASE_       )

          if not pairs:
            words.append(SCREAMING_SNAKE_CASE_       )
            continue

          while True:
            lowerCAmelCase_   :				int								= min(SCREAMING_SNAKE_CASE_    ,			key=lambda SCREAMING_SNAKE_CASE_       : self.bpe_ranks.get(SCREAMING_SNAKE_CASE_    ,			float('inf'       )       )       )
            if bigram not in self.bpe_ranks:
              break
            lowerCAmelCase_						,lowerCAmelCase_   :				Any								= bigram
            lowerCAmelCase_   :				int								= []
            lowerCAmelCase_   :				List[str]								= 0

            while i < len(SCREAMING_SNAKE_CASE_       ):
              try:
                lowerCAmelCase_   :				Optional[int]								= word.index(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
                new_word.extend(word[i:j]       )
                lowerCAmelCase_   :				List[Any]								= j
              except ValueError:
                new_word.extend(word[i:]       )
                break

              if word[i] == first and i < len(SCREAMING_SNAKE_CASE_       ) - 1 and word[i + 1] == second:
                new_word.append(first + second       )
                i += 2
              else:
                new_word.append(word[i]       )
                i += 1
            lowerCAmelCase_   :				Tuple								= tuple(SCREAMING_SNAKE_CASE_       )
            lowerCAmelCase_   :				Optional[int]								= new_word
            if len(SCREAMING_SNAKE_CASE_       ) == 1:
              break
            else:
              lowerCAmelCase_   :				int								= get_pairs(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				List[str]								= '@@ '.join(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Optional[int]								= word[:-4]

          lowerCAmelCase_   :				List[Any]								= word
          words.append(SCREAMING_SNAKE_CASE_       )
        return " ".join(SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				Optional[int]								= []

        lowerCAmelCase_   :				int								= re.findall(r'\S+\n?'    ,			SCREAMING_SNAKE_CASE_       )

        for token in words:
          split_tokens.extend(list(self.bpe(SCREAMING_SNAKE_CASE_       ).split(' '       )       )       )
        return split_tokens



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				List[str]								= token.lower()
        return self.encoder.get(SCREAMING_SNAKE_CASE_    ,			self.encoder.get(self.unk_token       )       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        return self.decoder.get(SCREAMING_SNAKE_CASE_    ,			self.unk_token       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        lowerCAmelCase_   :				Any								= ' '.join(SCREAMING_SNAKE_CASE_       ).replace('@@ '    ,			''       ).strip()
        return out_string





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      Optional[str] = None       ):
        if not os.path.isdir(SCREAMING_SNAKE_CASE_       ):
          logger.error(F"Vocabulary path ({save_directory}) should be a directory"       )
          return
        lowerCAmelCase_   :				List[str]								= os.path.join(
            SCREAMING_SNAKE_CASE_    ,			(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']       )
        lowerCAmelCase_   :				str								= os.path.join(
            SCREAMING_SNAKE_CASE_    ,			(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']       )

        with open(SCREAMING_SNAKE_CASE_    ,			'w'    ,			encoding='utf-8'       ) as f:
          f.write(json.dumps(self.encoder    ,			indent=2    ,			sort_keys=SCREAMING_SNAKE_CASE_    ,			ensure_ascii=SCREAMING_SNAKE_CASE_       ) + '\n'       )

        lowerCAmelCase_   :				Dict								= 0
        with open(SCREAMING_SNAKE_CASE_    ,			'w'    ,			encoding='utf-8'       ) as writer:
          writer.write('#version: 0.2\n'       )
          for bpe_tokens, token_index in sorted(self.bpe_ranks.items()    ,			key=lambda SCREAMING_SNAKE_CASE_       : kv[1]       ):
            if index != token_index:
              logger.warning(
                  F"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
                  ' Please check that the tokenizer is not corrupted!'       )
              lowerCAmelCase_   :				Tuple								= token_index
            writer.write(' '.join(SCREAMING_SNAKE_CASE_       ) + '\n'       )
            index += 1

        return vocab_file, merge_file










"""simple docstring"""







import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowercase__		:				List[Any]          =	logging.getLogger(__name__)

def 				UpperCamelCase_    (    )						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Tuple								= argparse.ArgumentParser(
      description='Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.'						)
  parser.add_argument(
      '--dataset_name'    ,   type=lowerCAmelCase__    ,   default='wikitext'    ,   help='Name of the training. Explore datasets at: hf.co/datasets.'    ,   )
  parser.add_argument(
      '--dataset_config'    ,   type=lowerCAmelCase__    ,   default='wikitext-103-raw-v1'    ,   help='Configuration name of the dataset.'						)
  parser.add_argument(
      '--tokenizer_name_or_path'    ,   type=lowerCAmelCase__    ,   default='sayakpaul/unigram-tokenizer-wikitext'    ,   help='Tokenizer identifier. Can be a local filepath or a Hub identifier.'    ,   )
  parser.add_argument(
      '--shard_size'    ,   type=lowerCAmelCase__    ,   default=1000    ,   help='Number of entries to go in a single shard.'    ,   )
  parser.add_argument('--split'    ,   type=lowerCAmelCase__    ,   default='train'    ,   choices=['train', 'test', 'validation']						)
  parser.add_argument(
      '--limit'    ,   default=lowerCAmelCase__    ,   type=lowerCAmelCase__    ,   help='Limit the number of shards (used for debugging).'    ,   )
  parser.add_argument(
      '--max_length'    ,   type=lowerCAmelCase__    ,   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=lowerCAmelCase__    ,   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.'    ,   )

  lowerCAmelCase_   :				List[Any]								= parser.parse_args()
  return args

def 				UpperCamelCase_    (    lowerCAmelCase__      :				int						)						->   Union[str, Any]:

  """simple docstring"""




  def fn(lowerCAmelCase__      :				Optional[Any]						):
    return tokenizer(examples['text']						)

  return fn

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Tuple						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				int								= []
  for i in range(len(tokenized_data['input_ids']						)						):
    lowerCAmelCase_   :				Tuple								= {
        '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]						)						),
    }
    lowerCAmelCase_   :				Union[str, Any]								= tf.train.Features(feature=lowerCAmelCase__						)
    lowerCAmelCase_   :				Dict								= tf.train.Example(features=lowerCAmelCase__						)
    lowerCAmelCase_   :				Optional[int]								= example.SerializeToString()
    records.append(lowerCAmelCase__						)
  return records

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[Any]						)						->   Optional[int]:

  """simple docstring"""




  lowerCAmelCase_   :				str								= datasets.load_dataset(args.dataset_name    ,   args.dataset_config    ,   split=args.split						)

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

  lowerCAmelCase_   :				List[str]								= 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						)
    lowerCAmelCase_   :				int								= os.path.join(args.output_dir    ,   args.split						)
    if not os.path.exists(lowerCAmelCase__						):
      os.makedirs(lowerCAmelCase__						)
  else:
    lowerCAmelCase_   :				Dict								= os.path.join(args.output_dir    ,   args.split						)

  # Tokenize the whole dataset at once.
  lowerCAmelCase_   :				Dict								= tokenize_function(lowerCAmelCase__						)
  lowerCAmelCase_   :				List[Any]								= dataset.map(lowerCAmelCase__    ,   batched=lowerCAmelCase__    ,   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(lowerCAmelCase__      :				List[Any]						):
    # Concatenate all texts.
    lowerCAmelCase_   :				int								= {k: sum(examples[k]    ,   []						) for k in examples.keys()}
    lowerCAmelCase_   :				Any								= 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 🫀
    lowerCAmelCase_   :				Union[str, Any]								= (total_length // args.max_length) * args.max_length
    # Split by chunks of max_len.
    lowerCAmelCase_   :				Optional[Any]								= {
        k: [t[i : i + args.max_length] for i in range(0    ,   lowerCAmelCase__    ,   args.max_length						)]
        for k, t in concatenated_examples.items()
    }
    return result

  lowerCAmelCase_   :				Optional[int]								= dataset_tokenized.map(lowerCAmelCase__    ,   batched=lowerCAmelCase__    ,   batch_size=1000    ,   num_proc=4						)

  lowerCAmelCase_   :				Optional[int]								= 0
  lowerCAmelCase_   :				Optional[Any]								= 0
  for shard in range(0    ,   len(lowerCAmelCase__						)    ,   args.shard_size						):
    lowerCAmelCase_   :				Dict								= grouped_dataset[shard : shard + args.shard_size]
    lowerCAmelCase_   :				Tuple								= len(dataset_snapshot['input_ids']						)
    lowerCAmelCase_   :				Optional[Any]								= os.path.join(lowerCAmelCase__    ,   f"dataset-{shard_count}-{records_containing}.tfrecord"						)
    lowerCAmelCase_   :				Tuple								= get_serialized_examples(lowerCAmelCase__						)

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

    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=lowerCAmelCase__						)


if __name__ == "__main__":
     lowercase__		:				int          =	parse_args()
     main(args)





