



"""simple docstring"""







import os
import re
import unicodedata
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import is_torch_available, logging


if is_torch_available():
     import torch


if TYPE_CHECKING:
     from transformers.pipelines.conversational import Conversation


lowercase__		:				List[str]          =	logging.get_logger(__name__)
lowercase__		:				Tuple          =	{"""vocab_file""": """spiece.model"""}

lowercase__		:				Optional[Any]          =	{
    """vocab_file""": {
        """AI-Sweden/gpt-sw3-126m""": """https://huggingface.co/AI-Sweden/gpt-sw3-126m/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-350m""": """https://huggingface.co/AI-Sweden/gpt-sw3-350m/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-1.6b""": """https://huggingface.co/AI-Sweden/gpt-sw3-1.6b/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-6.7b""": """https://huggingface.co/AI-Sweden/gpt-sw3-6.7b/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-20b""": """https://huggingface.co/AI-Sweden/gpt-sw3-20b/resolve/main/spiece.model""",
    }
}

lowercase__		:				Dict          =	{
    """AI-Sweden/gpt-sw3-126m""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-350m""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-1.6b""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-6.7b""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-20b""": 2_0_4_8,
}







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_				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      List[str]=False    ,			SCREAMING_SNAKE_CASE_				:      str=False    ,			SCREAMING_SNAKE_CASE_				:      Tuple=False    ,			SCREAMING_SNAKE_CASE_				:      str=None    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=None    ,			SCREAMING_SNAKE_CASE_				:      Dict=None    ,			SCREAMING_SNAKE_CASE_				:      Tuple=None    ,			SCREAMING_SNAKE_CASE_				:      Optional[Dict[str, Any]] = None    ,			**SCREAMING_SNAKE_CASE_				:      Any    ,			):
        lowerCAmelCase_   :				Any								= {} if sp_model_kwargs is None else sp_model_kwargs

        lowerCAmelCase_   :				Union[str, Any]								= kwargs.get('name_or_path'       )
        if name_or_path is None:
          logger.warning(
              'name_or_path not provided, will work for all GPTSw3 models except gpt-sw3-7b,'
              ' you are testing the model, this can safely be ignored'       )
          lowerCAmelCase_   :				List[str]								= 'None'

        # Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
        lowerCAmelCase_   :				Optional[Any]								= '<|endoftext|>' if eos_token is None else eos_token
        lowerCAmelCase_   :				List[str]								= '<unk>' if unk_token is None else unk_token
        if "gpt-sw3-7b" in name_or_path:
          lowerCAmelCase_   :				List[Any]								= unk_token if pad_token is None else pad_token
          lowerCAmelCase_   :				Dict								= eos_token if bos_token is None else bos_token
        else:
          lowerCAmelCase_   :				Optional[int]								= '<pad>' if pad_token is None else pad_token
          lowerCAmelCase_   :				List[str]								= '<s>' if bos_token is None else bos_token

        super().__init__(
            do_lower_case=SCREAMING_SNAKE_CASE_    ,			remove_space=SCREAMING_SNAKE_CASE_    ,			keep_accents=SCREAMING_SNAKE_CASE_    ,			bos_token=SCREAMING_SNAKE_CASE_    ,			eos_token=SCREAMING_SNAKE_CASE_    ,			unk_token=SCREAMING_SNAKE_CASE_    ,			pad_token=SCREAMING_SNAKE_CASE_    ,			sp_model_kwargs=self.sp_model_kwargs    ,			**SCREAMING_SNAKE_CASE_    ,			)

        lowerCAmelCase_   :				Optional[int]								= do_lower_case
        lowerCAmelCase_   :				Any								= remove_space
        lowerCAmelCase_   :				Tuple								= keep_accents
        lowerCAmelCase_   :				Optional[Any]								= vocab_file

        lowerCAmelCase_   :				List[Any]								= spm.SentencePieceProcessor(**self.sp_model_kwargs       )
        self.sp_model.Load(SCREAMING_SNAKE_CASE_       )

        # Used for whitespace normalization in input texts
        # fmt : off
        lowerCAmelCase_   :				Union[str, Any]								= {' ', ' ', ' ', ' ', ' ', '　', ' ', ' ', ' ', ' ', '￼', ''}
        # fmt : on

        # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
        lowerCAmelCase_   :				List[Any]								= re.compile(
            F"[{''.join(map(SCREAMING_SNAKE_CASE_    ,			list(range(0    ,			9       )       ) + list(range(1_1    ,			3_2       )       ) + list(range(1_2_7    ,			1_6_0       )       ) + [1_6_0, 1_7_3, 8_2_0_3]       )       )}]"       )



      def __getstate__(      self				:      Any       ):
        lowerCAmelCase_   :				Optional[Any]								= self.__dict__.copy()
        lowerCAmelCase_   :				int								= None
        return state



      def __setstate__(      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      Any       ):
        lowerCAmelCase_   :				Any								= d

        # for backward compatibility
        if not hasattr(self    ,			'sp_model_kwargs'       ):
          lowerCAmelCase_   :				List[Any]								= {}

        lowerCAmelCase_   :				Optional[Any]								= spm.SentencePieceProcessor(**self.sp_model_kwargs       )
        self.sp_model.Load(self.vocab_file       )



      @property
      # Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.vocab_size
      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        return len(self.sp_model       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        lowerCAmelCase_   :				int								= self.non_printing_characters_re.sub(''    ,			SCREAMING_SNAKE_CASE_       )

        # Normalize whitespaces
        lowerCAmelCase_   :				int								= ''.join([char if char not in self.whitespaces else ' ' for char in text]       )

        # NFC Unicode normalization
        lowerCAmelCase_   :				Union[str, Any]								= unicodedata.normalize('NFC'    ,			SCREAMING_SNAKE_CASE_       )
        return text



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			**SCREAMING_SNAKE_CASE_				:      Dict       ):
        lowerCAmelCase_   :				Optional[int]								= self.preprocess_text(SCREAMING_SNAKE_CASE_       )
        return self.sp_model.encode(SCREAMING_SNAKE_CASE_    ,			out_type=SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str       ):
        return self.sp_model.PieceToId(SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        return self.sp_model.IdToPiece(SCREAMING_SNAKE_CASE_       )



      @staticmethod
      def      SCREAMING_SNAKE_CASE__ (      SCREAMING_SNAKE_CASE_				:      str       ):
        return out_string



      def      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        lowerCAmelCase_   :				Optional[Any]								= []
        lowerCAmelCase_   :				Tuple								= ''
        lowerCAmelCase_   :				Tuple								= False
        for token in tokens:
          # make sure that special tokens are not decoded using sentencepiece model
          if token in self.all_special_tokens:
            # TODO: Check if this is needed, as it ensures that decode(encode(doc)) != doc by adding extra whitespace in the decoded document
            if not prev_is_special:
              out_string += " "

            out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_       ) + token
            lowerCAmelCase_   :				str								= True
            lowerCAmelCase_   :				List[str]								= []
          else:
            current_sub_tokens.append(SCREAMING_SNAKE_CASE_       )
            lowerCAmelCase_   :				Any								= False
        out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_       )

        return out_string



      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				Any								= {self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE_       ): i for i in range(self.vocab_size       )}
        vocab.update(self.added_tokens_encoder       )
        return vocab



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict    ,			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_   :				Optional[Any]								= os.path.join(
            SCREAMING_SNAKE_CASE_    ,			(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']       )

        if os.path.abspath(self.vocab_file       ) != os.path.abspath(SCREAMING_SNAKE_CASE_       ) and os.path.isfile(self.vocab_file       ):
          copyfile(self.vocab_file    ,			SCREAMING_SNAKE_CASE_       )
        elif not os.path.isfile(self.vocab_file       ):
          with open(SCREAMING_SNAKE_CASE_    ,			'wb'       ) as fi:
            lowerCAmelCase_   :				List[str]								= self.sp_model.serialized_model_proto()
            fi.write(SCREAMING_SNAKE_CASE_       )

        return (out_vocab_file,)



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      Union[str, List[str]]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, bool] = False       ):

        if isinstance(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Tuple								= self.preprocess_text(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Optional[int]								= self.sp_model.encode(SCREAMING_SNAKE_CASE_       )
        else:
          lowerCAmelCase_   :				Any								= [self.preprocess_text(SCREAMING_SNAKE_CASE_       ) for t in text]
          lowerCAmelCase_   :				List[str]								= self.sp_model.encode(SCREAMING_SNAKE_CASE_       )

        if return_tensors is True or return_tensors == "pt":
          lowerCAmelCase_   :				List[str]								= torch.tensor(SCREAMING_SNAKE_CASE_       )

        return token_ids



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Union[int, List[int]]       ):
        return self.sp_model.decode(SCREAMING_SNAKE_CASE_       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      "Conversation"       ):
        lowerCAmelCase_   :				str								= [F"User: {text}" if is_user else F"Bot: {text}" for is_user, text in conversation.iter_texts()]
        lowerCAmelCase_   :				Optional[int]								= (
            F"{self.eos_token}{self.bos_token}" + F"{self.bos_token}".join(SCREAMING_SNAKE_CASE_       ) + F"{self.bos_token}Bot:"
        )
        return self.encode(text=SCREAMING_SNAKE_CASE_       )










"""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)





