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


lowerCamelCase       	=						logging.get_logger(__name__)
lowerCamelCase       	=						{"""vocab_file""": """spiece.model"""}

lowerCamelCase       	=						{
    """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""",
    }
}

lowerCamelCase       	=						{
    """AI-Sweden/gpt-sw3-126m""": 2_048,
    """AI-Sweden/gpt-sw3-350m""": 2_048,
    """AI-Sweden/gpt-sw3-1.6b""": 2_048,
    """AI-Sweden/gpt-sw3-6.7b""": 2_048,
    """AI-Sweden/gpt-sw3-20b""": 2_048,
}
class 				lowercase__       (   __lowercase				):







       '''simple docstring'''


       UpperCamelCase						=					VOCAB_FILES_NAMES
       UpperCamelCase						=					PRETRAINED_VOCAB_FILES_MAP
       UpperCamelCase						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCamelCase						=					['''input_ids''', '''attention_mask''']



       def __init__(							self      :  Optional[int]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  int=False		,    _UpperCAmelCase      :  List[str]=False		,    _UpperCAmelCase      :  Tuple=False		,    _UpperCAmelCase      :  Dict=None		,    _UpperCAmelCase      :  Tuple=None		,    _UpperCAmelCase      :  Optional[int]=None		,    _UpperCAmelCase      :  int=None		,    _UpperCAmelCase      :  Optional[Dict[str, Any]] = None		,    **_UpperCAmelCase      :  List[Any]		,    )			->						None:


             '''simple docstring'''




             UpperCAmelCase_      			=							{} if sp_model_kwargs is None else sp_model_kwargs

             UpperCAmelCase_      			=							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"   )
                   UpperCAmelCase_      			=							"None"

             # Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
             UpperCAmelCase_      			=							"<|endoftext|>" if eos_token is None else eos_token
             UpperCAmelCase_      			=							"<unk>" if unk_token is None else unk_token
             if "gpt-sw3-7b" in name_or_path:
                   UpperCAmelCase_      			=							unk_token if pad_token is None else pad_token
                   UpperCAmelCase_      			=							eos_token if bos_token is None else bos_token
             else:
                   UpperCAmelCase_      			=							"<pad>" if pad_token is None else pad_token
                   UpperCAmelCase_      			=							"<s>" if bos_token is None else bos_token

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

             UpperCAmelCase_      			=							do_lower_case
             UpperCAmelCase_      			=							remove_space
             UpperCAmelCase_      			=							keep_accents
             UpperCAmelCase_      			=							vocab_file

             UpperCAmelCase_      			=							spm.SentencePieceProcessor(**self.sp_model_kwargs   )
             self.sp_model.Load(UpperCAmelCase__   )

             # Used for whitespace normalization in input texts
             # fmt : off
             UpperCAmelCase_      			=							{" ", " ", " ", " ", " ", "　", " ", " ", " ", " ", "￼", ""}
             # fmt : on

             # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
             UpperCAmelCase_      			=							re.compile(
                 F"""[{''.join(map(UpperCAmelCase__		,    list(range(0		,    9   )   ) + list(range(11		,    32   )   ) + list(range(127		,    160   )   ) + [160, 173, 8203]   )   )}]"""   )



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


             '''simple docstring'''




             UpperCAmelCase_      			=							self.__dict__.copy()
             UpperCAmelCase_      			=							None
             return state



       def __setstate__(							self      :  Tuple		,    _UpperCAmelCase      :  Union[str, Any]   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							d

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

             UpperCAmelCase_      			=							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 					lowercase__    (							self      :  Any   )			->						int:


             '''simple docstring'''




             return len(self.sp_model   )



       def 					lowercase__    (							self      :  Any		,    _UpperCAmelCase      :  str   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.non_printing_characters_re.sub(""		,    UpperCAmelCase__   )

             # Normalize whitespaces
             UpperCAmelCase_      			=							"".join([char if char not in self.whitespaces else " " for char in text]   )

             # NFC Unicode normalization
             UpperCAmelCase_      			=							unicodedata.normalize("NFC"		,    UpperCAmelCase__   )
             return text



       def 					lowercase__    (							self      :  Dict		,    _UpperCAmelCase      :  str		,    **_UpperCAmelCase      :  Optional[int]   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.preprocess_text(UpperCAmelCase__   )
             return self.sp_model.encode(UpperCAmelCase__		,    out_type=UpperCAmelCase__   )



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  str   )			->						int:


             '''simple docstring'''




             return self.sp_model.PieceToId(UpperCAmelCase__   )



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  int   )			->						str:


             '''simple docstring'''




             return self.sp_model.IdToPiece(UpperCAmelCase__   )



       @staticmethod
       def 					lowercase__    (							_UpperCAmelCase      :  str   )			->						str:


             '''simple docstring'''




             return out_string



       def 					lowercase__    (							self      :  Any		,    _UpperCAmelCase      :  List[str]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							[]
             UpperCAmelCase_      			=							""
             UpperCAmelCase_      			=							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(UpperCAmelCase__   ) + token
                         UpperCAmelCase_      			=							True
                         UpperCAmelCase_      			=							[]
                   else:
                         current_sub_tokens.append(UpperCAmelCase__   )
                         UpperCAmelCase_      			=							False
             out_string += self.sp_model.decode(UpperCAmelCase__   )

             return out_string



       def 					lowercase__    (							self      :  int   )			->						Dict[str, int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							{self.convert_ids_to_tokens(UpperCAmelCase__   ): i for i in range(self.vocab_size   )}
             vocab.update(self.added_tokens_encoder   )
             return vocab



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Optional[str] = None   )			->						Tuple[str]:


             '''simple docstring'''




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

             if os.path.abspath(self.vocab_file   ) != os.path.abspath(UpperCAmelCase__   ) and os.path.isfile(self.vocab_file   ):
                   copyfile(self.vocab_file		,    UpperCAmelCase__   )
             elif not os.path.isfile(self.vocab_file   ):
                   with open(UpperCAmelCase__		,    "wb"   ) as fi:
                         UpperCAmelCase_      			=							self.sp_model.serialized_model_proto()
                         fi.write(UpperCAmelCase__   )

             return (out_vocab_file,)



       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  Union[str, List[str]]		,    _UpperCAmelCase      :  Union[str, bool] = False   )			->						Union[List[int], List[List[int]], "torch.Tensor"]:


             '''simple docstring'''




             if isinstance(UpperCAmelCase__		,    UpperCAmelCase__   ):
                   UpperCAmelCase_      			=							self.preprocess_text(UpperCAmelCase__   )
                   UpperCAmelCase_      			=							self.sp_model.encode(UpperCAmelCase__   )
             else:
                   UpperCAmelCase_      			=							[self.preprocess_text(UpperCAmelCase__   ) for t in text]
                   UpperCAmelCase_      			=							self.sp_model.encode(UpperCAmelCase__   )

             if return_tensors is True or return_tensors == "pt":
                   UpperCAmelCase_      			=							torch.tensor(UpperCAmelCase__   )

             return token_ids



       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  Union[int, List[int]]   )			->						str:


             '''simple docstring'''




             return self.sp_model.decode(UpperCAmelCase__   )







       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  "Conversation"   )			->						List[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							[F"""User: {text}""" if is_user else F"""Bot: {text}""" for is_user, text in conversation.iter_texts()]
             UpperCAmelCase_      			=							(
                 F"""{self.eos_token}{self.bos_token}""" + F"""{self.bos_token}""".join(UpperCAmelCase__   ) + F"""{self.bos_token}Bot:"""
             )
             return self.encode(text=UpperCAmelCase__   )









"""simple docstring"""





import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


lowerCamelCase       	=						datasets.utils.logging.get_logger(__name__)
@dataclass
class 				lowercase__       (   datasets.BuilderConfig				):







       '''simple docstring'''


       UpperCamelCase						=					None
       UpperCamelCase						=					"utf-8"
       UpperCamelCase						=					None
       UpperCamelCase						=					None
       UpperCamelCase						=					True  # deprecated
       UpperCamelCase						=					None  # deprecated
       UpperCamelCase						=					10 << 20  # 10MB
       UpperCamelCase						=					None



class 				lowercase__       (   datasets.ArrowBasedBuilder				):







       '''simple docstring'''


       UpperCamelCase						=					JsonConfig



       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             if self.config.block_size is not None:
                   logger.warning("The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead"   )
                   UpperCAmelCase_      			=							self.config.block_size
             if self.config.use_threads is not True:
                   logger.warning(
                       "The JSON loader parameter `use_threads` is deprecated and doesn't have any effect anymore."   )
             if self.config.newlines_in_values is not None:
                   raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported"   )
             return datasets.DatasetInfo(features=self.config.features   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  Optional[int]   )			->						Optional[int]:


             '''simple docstring'''




             if not self.config.data_files:
                   raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""   )
             UpperCAmelCase_      			=							dl_manager.download_and_extract(self.config.data_files   )
             if isinstance(_UpperCAmelCase		,    (str, list, tuple)   ):
                   UpperCAmelCase_      			=							data_files
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,    gen_kwargs={"files": files}   )]
             UpperCAmelCase_      			=							[]
             for split_name, files in data_files.items():
                   if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):
                         UpperCAmelCase_      			=							[files]
                   UpperCAmelCase_      			=							[dl_manager.iter_files(_UpperCAmelCase   ) for file in files]
                   splits.append(datasets.SplitGenerator(name=_UpperCAmelCase		,    gen_kwargs={"files": files}   )   )
             return splits



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  pa.Table   )			->						pa.Table:


             '''simple docstring'''




             if self.config.features is not None:
                   # adding missing columns
                   for column_name in set(self.config.features   ) - set(pa_table.column_names   ):
                         UpperCAmelCase_      			=							self.config.features.arrow_schema.field(_UpperCAmelCase   ).type
                         UpperCAmelCase_      			=							pa_table.append_column(_UpperCAmelCase		,    pa.array([None] * len(_UpperCAmelCase   )		,    type=_UpperCAmelCase   )   )
                   # more expensive cast to support nested structures with keys in a different order
                   # allows str <-> int/float or str to Audio for example
                   UpperCAmelCase_      			=							table_cast(_UpperCAmelCase		,    self.config.features.arrow_schema   )
             return pa_table







       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  List[Any]   )			->						str:


             '''simple docstring'''




             for file_idx, file in enumerate(itertools.chain.from_iterable(_UpperCAmelCase   )   ):
                   # If the file is one json object and if we need to look at the list of items in one specific field
                   if self.config.field is not None:
                         with open(_UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                               UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )

                         # We keep only the field we are interested in
                         UpperCAmelCase_      			=							dataset[self.config.field]

                         # We accept two format: a list of dicts or a dict of lists
                         if isinstance(_UpperCAmelCase		,    (list, tuple)   ):
                               UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                               UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                         else:
                               UpperCAmelCase_      			=							dataset
                         UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                         yield file_idx, self._cast_table(_UpperCAmelCase   )

                   # If the file has one json object per line
                   else:
                         with open(_UpperCAmelCase		,    "rb"   ) as f:
                               UpperCAmelCase_      			=							0
                               # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                               # Set a default minimum value of 16kB if the chunk size is really small
                               UpperCAmelCase_      			=							max(self.config.chunksize // 32		,    16 << 10   )
                               UpperCAmelCase_      			=							(
                                   self.config.encoding_errors if self.config.encoding_errors is not None else "strict"
                               )
                               while True:
                                     UpperCAmelCase_      			=							f.read(self.config.chunksize   )
                                     if not batch:
                                           break
                                     # Finish current line
                                     try:
                                           batch += f.readline()
                                     except (AttributeError, io.UnsupportedOperation):
                                           batch += readline(_UpperCAmelCase   )
                                     # PyArrow only accepts utf-8 encoded bytes
                                     if self.config.encoding != "utf-8":
                                           UpperCAmelCase_      			=							batch.decode(self.config.encoding		,    errors=_UpperCAmelCase   ).encode("utf-8"   )
                                     try:
                                           while True:
                                                 try:
                                                       UpperCAmelCase_      			=							paj.read_json(
                                                           io.BytesIO(_UpperCAmelCase   )		,    read_options=paj.ReadOptions(block_size=_UpperCAmelCase   )   )
                                                       break
                                                 except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                                       if (
                                                           isinstance(_UpperCAmelCase		,    pa.ArrowInvalid   )
                                                           and "straddling" not in str(_UpperCAmelCase   )
                                                           or block_size > len(_UpperCAmelCase   )
                                                       ):
                                                             raise
                                                       else:
                                                             # Increase the block size in case it was too small.
                                                             # The block size will be reset for the next file.
                                                             logger.debug(
                                                                 F"""Batch of {len(_UpperCAmelCase   )} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."""   )
                                                             block_size *= 2
                                     except pa.ArrowInvalid as e:
                                           try:
                                                 with open(
                                                     _UpperCAmelCase		,    encoding=self.config.encoding		,    errors=self.config.encoding_errors   ) as f:
                                                       UpperCAmelCase_      			=							json.load(_UpperCAmelCase   )
                                           except json.JSONDecodeError:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise e
                                           # If possible, parse the file as a list of json objects and exit the loop
                                           if isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ):  # list is the only sequence type supported in JSON
                                                 try:
                                                       UpperCAmelCase_      			=							set().union(*[row.keys() for row in dataset]   )
                                                       UpperCAmelCase_      			=							{col: [row.get(_UpperCAmelCase   ) for row in dataset] for col in keys}
                                                       UpperCAmelCase_      			=							pa.Table.from_pydict(_UpperCAmelCase   )
                                                 except (pa.ArrowInvalid, AttributeError) as e:
                                                       logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                       raise ValueError(F"""Not able to read records in the JSON file at {file}."""   ) from None
                                                 yield file_idx, self._cast_table(_UpperCAmelCase   )
                                                 break
                                           else:
                                                 logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase   )}: {e}"""   )
                                                 raise ValueError(
                                                     F"""Not able to read records in the JSON file at {file}. """
                                                     F"""You should probably indicate the field of the JSON file containing your records. """
                                                     F"""This JSON file contain the following fields: {str(list(dataset.keys()   )   )}. """
                                                     F"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """   ) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                     yield (file_idx, batch_idx), self._cast_table(_UpperCAmelCase   )
                                     batch_idx += 1




