



"""simple docstring"""
import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


A:   List[str]								=		datasets.utils.logging.get_logger(__name__)



@dataclass
class       SCREAMING_SNAKE_CASE__  (					datasets.BuilderConfig ):
   __lowerCAmelCase  :   str						=						10_000
   __lowerCAmelCase  :   List[str]						=						None
   __lowerCAmelCase  :   Union[str, Any]						=						None



class       SCREAMING_SNAKE_CASE__  (					datasets.ArrowBasedBuilder ):
   __lowerCAmelCase  :   Union[str, Any]						=						ParquetConfig

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    return datasets.DatasetInfo(features=self.config.features						)

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      Tuple:





    '''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	:					List[str]	=	dl_manager.download_and_extract(self.config.data_files						)
    if isinstance(__lowerCAmelCase      ,     (str, list, tuple)						):
     UpperCAmelCase	:					Any	=	data_files
     if isinstance(__lowerCAmelCase      ,     __lowerCAmelCase						):
      UpperCAmelCase	:					Optional[Any]	=	[files]
     # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
     UpperCAmelCase	:					Dict	=	[dl_manager.iter_files(__lowerCAmelCase						) for file in files]
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,     gen_kwargs={"""files""": files}						)]
    UpperCAmelCase	:					Any	=	[]
    for split_name, files in data_files.items():
     if isinstance(__lowerCAmelCase      ,     __lowerCAmelCase						):
      UpperCAmelCase	:					Dict	=	[files]
     # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
     UpperCAmelCase	:					Optional[Any]	=	[dl_manager.iter_files(__lowerCAmelCase						) for file in files]
     # Infer features is they are stoed in the arrow schema
     if self.info.features is None:
      for file in itertools.chain.from_iterable(__lowerCAmelCase						):
       with open(__lowerCAmelCase      ,     """rb"""						) as f:
        UpperCAmelCase	:					Union[str, Any]	=	datasets.Features.from_arrow_schema(pq.read_schema(__lowerCAmelCase						)						)
       break
     splits.append(datasets.SplitGenerator(name=__lowerCAmelCase      ,     gen_kwargs={"""files""": files}						)						)
    return splits

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      str:





    '''simple docstring'''




    if self.info.features is not None:
     # more expensive cast to support nested features with keys in a different order
     # allows str <-> int/float or str to Audio for example
     UpperCAmelCase	:					str	=	table_cast(__lowerCAmelCase      ,     self.info.features.arrow_schema						)
    return pa_table




   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      str:





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	self.info.features.arrow_schema if self.info.features is not None else None
    if self.info.features is not None and self.config.columns is not None:
     if sorted(field.name for field in schema						) != sorted(self.config.columns						):
      raise ValueError(
          F"Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\'"						)
    for file_idx, file in enumerate(itertools.chain.from_iterable(__lowerCAmelCase						)						):
     with open(__lowerCAmelCase      ,     """rb"""						) as f:
      UpperCAmelCase	:					Tuple	=	pq.ParquetFile(__lowerCAmelCase						)
      try:
       for batch_idx, record_batch in enumerate(
           parquet_file.iter_batches(batch_size=self.config.batch_size      ,     columns=self.config.columns						)						):
        UpperCAmelCase	:					Tuple	=	pa.Table.from_batches([record_batch]						)
        # 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 F"{file_idx}_{batch_idx}", self._cast_table(__lowerCAmelCase						)
      except ValueError as e:
       logger.error(F"Failed to read file \'{file}\' with error {type(__lowerCAmelCase						)}: {e}"						)
       raise







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


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

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

_a 						=      {
    "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 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







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





			def __init__(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase=None       ,  __lowerCAmelCase = None       ,  **__lowerCAmelCase       ,  ):



				'''simple docstring'''







				lowerCamelCase__								=		{} if sp_model_kwargs is None else sp_model_kwargs

				lowerCamelCase__								=		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__								=		'''None'''

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

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

				lowerCamelCase__								=		do_lower_case
				lowerCamelCase__								=		remove_space
				lowerCamelCase__								=		keep_accents
				lowerCamelCase__								=		vocab_file

				lowerCamelCase__								=		spm.SentencePieceProcessor(**self.sp_model_kwargs							)
				self.sp_model.Load(__lowerCAmelCase							)

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

				# Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
				lowerCamelCase__								=		re.compile(
				    F'[{"".join(map(__lowerCAmelCase       ,  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							):



				'''simple docstring'''







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





			def __setstate__(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		d

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

				lowerCamelCase__								=		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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return len(self.sp_model							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		self.non_printing_characters_re.sub(''''''       ,  __lowerCAmelCase							)

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

				# NFC Unicode normalization
				lowerCamelCase__								=		unicodedata.normalize('''NFC'''       ,  __lowerCAmelCase							)
				return text





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  **__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__								=		self.preprocess_text(__lowerCAmelCase							)
				return self.sp_model.encode(__lowerCAmelCase       ,  out_type=__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				return self.sp_model.PieceToId(__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				return self.sp_model.IdToPiece(__lowerCAmelCase							)





			@staticmethod
			def 		__lowerCamelCase				(    __lowerCAmelCase							):



				'''simple docstring'''







				return out_string





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







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

				return out_string





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase = None							):



				'''simple docstring'''







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

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

				return (out_vocab_file,)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase = False							):



				'''simple docstring'''







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

				if return_tensors is True or return_tensors == "pt":
					lowerCamelCase__								=		torch.tensor(__lowerCAmelCase							)

				return token_ids





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				return self.sp_model.decode(__lowerCAmelCase							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







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



