


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 : Dict		=				logging.get_logger(__name__)
a : List[Any]		=				{'''vocab_file''': '''spiece.model'''}

a : List[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''',
    }
}

a : List[Any]		=				{
    '''AI-Sweden/gpt-sw3-126m''': 2048,
    '''AI-Sweden/gpt-sw3-350m''': 2048,
    '''AI-Sweden/gpt-sw3-1.6b''': 2048,
    '''AI-Sweden/gpt-sw3-6.7b''': 2048,
    '''AI-Sweden/gpt-sw3-20b''': 2048,
}
class 	lowerCamelCase_       (				lowerCAmelCase__				):



					'''simple docstring'''







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






					def __init__(				self ,       snake_case_ ,       snake_case_=False ,       snake_case_=False ,       snake_case_=False ,       snake_case_=None ,       snake_case_=None ,       snake_case_=None ,       snake_case_=None ,       snake_case_ = None ,       **snake_case_ ,       )   ->			None:







								'''simple docstring'''

								__lowercase 					=   {} if sp_model_kwargs is None else sp_model_kwargs

								__lowercase 					=   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'''						)
											__lowercase 					=   '''None'''

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

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

								__lowercase 					=   do_lower_case
								__lowercase 					=   remove_space
								__lowercase 					=   keep_accents
								__lowercase 					=   vocab_file

								__lowercase 					=   spm.SentencePieceProcessor(**self.sp_model_kwargs						)
								self.sp_model.Load(snake_case_						)

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

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







								'''simple docstring'''

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






					def __setstate__(				self ,       snake_case_						)   ->			str:







								'''simple docstring'''

								__lowercase 					=   d

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

								__lowercase 					=   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 A      (				self						)   ->			int:







								'''simple docstring'''

								return len(self.sp_model						)






					def A      (				self ,       snake_case_						)   ->			str:







								'''simple docstring'''

								__lowercase 					=   self.non_printing_characters_re.sub('''''' ,       snake_case_						)

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

								# NFC Unicode normalization
								__lowercase 					=   unicodedata.normalize('''NFC''' ,       snake_case_						)
								return text






					def A      (				self ,       snake_case_ ,       **snake_case_						)   ->			List[str]:







								'''simple docstring'''

								__lowercase 					=   self.preprocess_text(snake_case_						)
								return self.sp_model.encode(snake_case_ ,       out_type=snake_case_						)






					def A      (				self ,       snake_case_						)   ->			int:







								'''simple docstring'''

								return self.sp_model.PieceToId(snake_case_						)






					def A      (				self ,       snake_case_						)   ->			str:







								'''simple docstring'''

								return self.sp_model.IdToPiece(snake_case_						)






					@staticmethod
					def A      (				snake_case_						)   ->			str:







								'''simple docstring'''

								return out_string






					def A      (				self ,       snake_case_						)   ->			str:







								'''simple docstring'''

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

								return out_string






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







								'''simple docstring'''

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






					def A      (				self ,       snake_case_ ,       snake_case_ = None						)   ->			Tuple[str]:







								'''simple docstring'''

								if not os.path.isdir(snake_case_						):
											logger.error(F'Vocabulary path ({save_directory}) should be a directory'						)
											return
								__lowercase 					=   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_						) and os.path.isfile(self.vocab_file						):
											copyfile(self.vocab_file ,       snake_case_						)
								elif not os.path.isfile(self.vocab_file						):
											with open(snake_case_ ,       '''wb'''						) as fi:
														__lowercase 					=   self.sp_model.serialized_model_proto()
														fi.write(snake_case_						)

								return (out_vocab_file,)






					def A      (				self ,       snake_case_ ,       snake_case_ = False						)   ->			Union[List[int], List[List[int]], "torch.Tensor"]:







								'''simple docstring'''

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

								if return_tensors is True or return_tensors == "pt":
											__lowercase 					=   torch.tensor(snake_case_						)

								return token_ids






					def A      (				self ,       snake_case_						)   ->			str:







								'''simple docstring'''

								return self.sp_model.decode(snake_case_						)






					def A      (				self ,       snake_case_						)   ->			List[int]:







								'''simple docstring'''

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






import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available


def 		lowercase_  (     *_UpperCamelCase  ):


			'''simple docstring'''







			if not isinstance(_UpperCamelCase       ,			_UpperCamelCase  ):
						__lowercase 					=   list(_UpperCamelCase  )
			for i in range(len(_UpperCamelCase  )  ):
						__lowercase 					=   None
			gc.collect()
			if is_xpu_available():
						torch.xpu.empty_cache()
			elif is_npu_available():
						torch.npu.empty_cache()
			else:
						torch.cuda.empty_cache()
			return objects


def 		lowercase_  (     _UpperCamelCase  ):


			'''simple docstring'''







			__lowercase 					=   [
			    '''CUDA out of memory.''',  # CUDA OOM
			    '''cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.''',  # CUDNN SNAFU
			    '''DefaultCPUAllocator: can\'t allocate memory''',  # CPU OOM
			]
			if isinstance(_UpperCamelCase       ,			_UpperCamelCase  ) and len(exception.args  ) == 1:
						return any(err in exception.args[0] for err in _statements  )
			return False





def 		lowercase_  (     _UpperCamelCase = None       ,			_UpperCamelCase = 1_28  ):


			'''simple docstring'''







			if function is None:
						return functools.partial(_UpperCamelCase       ,			starting_batch_size=_UpperCamelCase  )

			__lowercase 					=   starting_batch_size

			def decorator(*_UpperCamelCase       ,			**_UpperCamelCase  ):
						nonlocal batch_size
						gc.collect()
						if is_xpu_available():
									torch.xpu.empty_cache()
						elif is_npu_available():
									torch.npu.empty_cache()
						else:
									torch.cuda.empty_cache()
						__lowercase 					=   list(inspect.signature(_UpperCamelCase  ).parameters.keys()  )
						# Guard against user error
						if len(_UpperCamelCase  ) < (len(_UpperCamelCase  ) + 1):
									__lowercase 					=   ''', '''.join([F'{arg}={value}' for arg, value in zip(params[1:]       ,			args[1:]  )]  )
									raise TypeError(
									    F'Batch size was passed into `{function.__name__}` as the first argument when called.'
									    F'Remove this as the decorator already does so: `{function.__name__}({arg_str})`'  )
						while True:
									if batch_size == 0:
												raise RuntimeError('''No executable batch size found, reached zero.'''  )
									try:
												return function(_UpperCamelCase       ,			*_UpperCamelCase       ,			**_UpperCamelCase  )
									except Exception as e:
												if should_reduce_batch_size(_UpperCamelCase  ):
															gc.collect()
															if is_xpu_available():
																		torch.xpu.empty_cache()
															elif is_npu_available():
																		torch.npu.empty_cache()
															else:
																		torch.cuda.empty_cache()
															batch_size //= 2
												else:
															raise

			return decorator


