
import enum
import warnings

from .. import MODEL_FOR_CAUSAL_LM_MAPPING, TF_MODEL_FOR_CAUSAL_LM_MAPPING
from ..utils import add_end_docstrings, is_tf_available
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
					import tensorflow as tf





class 							UpperCAmelCase      (  enum.Enum  ):




	'''simple docstring'''





	snake_case_					    =     0
	snake_case_					    =     1
	snake_case_					    =     2





@add_end_docstrings(__SCREAMING_SNAKE_CASE  )
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	snake_case_					    =     "\n    In 1991, the remains of Russian Tsar Nicholas II and his family (except for Alexei and Maria) are discovered. The\n    voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the remainder of the story. 1883 Western\n    Siberia, a young Grigori Rasputin is asked by his father and a group of men to perform magic. Rasputin has a vision\n    and denounces one of the men as a horse thief. Although his father initially slaps him for making such an\n    accusation, Rasputin watches as the man is chased outside and beaten. Twenty years later, Rasputin sees a vision of\n    the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous, with people, even a bishop,\n    begging for his blessing. <eod> </s> <eos>\n    "






	def __init__(       self					: List[str]							,*A					: Optional[Any]							,**A					: str ):
					super().__init__(*A							,**A )
					self.check_model_type(
					    TF_MODEL_FOR_CAUSAL_LM_MAPPING if self.framework == "tf" else MODEL_FOR_CAUSAL_LM_MAPPING )
					if "prefix" not in self._preprocess_params:
									# This is very specific. The logic is quite complex and needs to be done
									# as a "default".
									# It also defines both some preprocess_kwargs and generate_kwargs
									# which is why we cannot put them in their respective methods.
									__A   			= None
									if self.model.config.prefix is not None:
													__A   			= self.model.config.prefix
									if prefix is None and self.model.__class__.__name__ in [
									    "XLNetLMHeadModel",
									    "TransfoXLLMHeadModel",
									    "TFXLNetLMHeadModel",
									    "TFTransfoXLLMHeadModel",
									]:
													# For XLNet and TransformerXL we add an article to the prompt to give more state to the model.
													__A   			= self.XL_PREFIX
									if prefix is not None:
													# Recalculate some generate_kwargs linked to prefix.
													__A			,							__A			,							__A   			= self._sanitize_parameters(prefix=A							,**self._forward_params )
													__A   			= {**self._preprocess_params, **preprocess_params}
													__A   			= {**self._forward_params, **forward_params}






	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: Union[str, Any]=None							,A					: int=None							,A					: Union[str, Any]=None							,A					: Optional[Any]=None							,A					: Tuple=None							,A					: int=None							,A					: str=None							,A					: Dict=None							,**A					: Optional[Any]							,):
					__A   			= {}
					if prefix is not None:
									__A   			= prefix
					if prefix:
									__A   			= self.tokenizer(
									    A							,padding=A							,add_special_tokens=A							,return_tensors=self.framework )
									__A   			= prefix_inputs["input_ids"].shape[-1]

					if handle_long_generation is not None:
									if handle_long_generation not in {"hole"}:
													raise ValueError(
													    f'''{handle_long_generation} is not a valid value for `handle_long_generation` parameter expected'''
													    " [None, 'hole']" )
									__A   			= handle_long_generation

					preprocess_params.update(A )
					__A   			= generate_kwargs

					__A   			= {}
					if return_full_text is not None and return_type is None:
									if return_text is not None:
													raise ValueError("`return_text` is mutually exclusive with `return_full_text`" )
									if return_tensors is not None:
													raise ValueError("`return_full_text` is mutually exclusive with `return_tensors`" )
									__A   			= ReturnType.FULL_TEXT if return_full_text else ReturnType.NEW_TEXT
					if return_tensors is not None and return_type is None:
									if return_text is not None:
													raise ValueError("`return_text` is mutually exclusive with `return_tensors`" )
									__A   			= ReturnType.TENSORS
					if return_type is not None:
									__A   			= return_type
					if clean_up_tokenization_spaces is not None:
									__A   			= clean_up_tokenization_spaces

					if stop_sequence is not None:
									__A   			= self.tokenizer.encode(A							,add_special_tokens=A )
									if len(A ) > 1:
													warnings.warn(
													    "Stopping on a multiple token sequence is not yet supported on transformers. The first token of"
													    " the stop sequence will be used as the stop sequence string in the interim." )
									__A   			= stop_sequence_ids[0]

					return preprocess_params, forward_params, postprocess_params






	def 				UpperCamelCase_  (       self					: Union[str, Any]							,*A					: Any							,**A					: int ):
					# Parse arguments
					if self.model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
									kwargs.update({"add_space_before_punct_symbol": True} )

					return super()._parse_and_tokenize(*A							,**A )






	def __call__(       self					: List[str]							,A					: int							,**A					: Dict ):
					return super().__call__(A							,**A )






	def 				UpperCamelCase_  (       self					: str							,A					: Optional[int]							,A					: Any=""							,A					: Union[str, Any]=None							,**A					: Any ):
					__A   			= self.tokenizer(
					    prefix + prompt_text							,padding=A							,add_special_tokens=A							,return_tensors=self.framework )
					__A   			= prompt_text

					if handle_long_generation == "hole":
									__A   			= inputs["input_ids"].shape[-1]
									if "max_new_tokens" in generate_kwargs:
													__A   			= generate_kwargs["max_new_tokens"]
									else:
													__A   			= generate_kwargs.get("max_length"							,self.model.config.max_length ) - cur_len
													if new_tokens < 0:
																	raise ValueError("We cannot infer how many new tokens are expected" )
									if cur_len + new_tokens > self.tokenizer.model_max_length:
													__A   			= self.tokenizer.model_max_length - new_tokens
													if keep_length <= 0:
																	raise ValueError(
																	    "We cannot use `hole` to handle this generation the number of desired tokens exceeds the"
																	    " models max length" )

													__A   			= inputs["input_ids"][:, -keep_length:]
													if "attention_mask" in inputs:
																	__A   			= inputs["attention_mask"][:, -keep_length:]

					return inputs






	def 				UpperCamelCase_  (       self					: int							,A					: str							,**A					: List[Any] ):
					__A   			= model_inputs["input_ids"]
					__A   			= model_inputs.get("attention_mask"							,A )
					# Allow empty prompts
					if input_ids.shape[1] == 0:
									__A   			= None
									__A   			= None
									__A   			= 1
					else:
									__A   			= input_ids.shape[0]
					__A   			= model_inputs.pop("prompt_text" )

					# If there is a prefix, we may need to adjust the generation length. Do so without permanently modifying
					# generate_kwargs, as some of the parameterization may come from the initialization of the pipeline.
					__A   			= generate_kwargs.pop("prefix_length"							,0 )
					if prefix_length > 0:
									__A   			= "max_new_tokens" in generate_kwargs or (
									    "generation_config" in generate_kwargs
									    and generate_kwargs["generation_config"].max_new_tokens is not None
									)
									if not has_max_new_tokens:
													__A   			= generate_kwargs.get("max_length" ) or self.model.config.max_length
													generate_kwargs["max_length"] += prefix_length
									__A   			= "min_new_tokens" in generate_kwargs or (
									    "generation_config" in generate_kwargs
									    and generate_kwargs["generation_config"].min_new_tokens is not None
									)
									if not has_min_new_tokens and "min_length" in generate_kwargs:
													generate_kwargs["min_length"] += prefix_length

        # BS x SL
					__A   			= self.model.generate(input_ids=A							,attention_mask=A							,**A )
					__A   			= generated_sequence.shape[0]
					if self.framework == "pt":
									__A   			= generated_sequence.reshape(A							,out_b // in_b							,*generated_sequence.shape[1:] )
					elif self.framework == "tf":
									__A   			= tf.reshape(A							,(in_b, out_b // in_b, *generated_sequence.shape[1:]) )
					return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text}







	def 				UpperCamelCase_  (       self					: List[Any]							,A					: Dict							,A					: Optional[int]=ReturnType.FULL_TEXT							,A					: Dict=True ):
					__A   			= model_outputs["generated_sequence"][0]
					__A   			= model_outputs["input_ids"]
					__A   			= model_outputs["prompt_text"]
					__A   			= generated_sequence.numpy().tolist()
					__A   			= []
					for sequence in generated_sequence:
									if return_type == ReturnType.TENSORS:
													__A   			= {"generated_token_ids": sequence}
									elif return_type in {ReturnType.NEW_TEXT, ReturnType.FULL_TEXT}:
													# Decode text
													__A   			= self.tokenizer.decode(
													    A							,skip_special_tokens=A							,clean_up_tokenization_spaces=A							,)

													# Remove PADDING prompt of the sequence if XLNet or Transfo-XL model is used
													if input_ids is None:
																	__A   			= 0
													else:
																	__A   			= len(
																	    self.tokenizer.decode(
																	        input_ids[0]							,skip_special_tokens=A							,clean_up_tokenization_spaces=A							,) )

													if return_type == ReturnType.FULL_TEXT:
																	__A   			= prompt_text + text[prompt_length:]
													else:
																	__A   			= text[prompt_length:]

													__A   			= {"generated_text": all_text}
									records.append(A )

					return records


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


SCREAMING_SNAKE_CASE							:List[str]							 =				logging.get_logger(__name__)
SCREAMING_SNAKE_CASE							:List[str]							 =				{'vocab_file': 'spiece.model'}

SCREAMING_SNAKE_CASE							:Dict							 =				{
    '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',
    }
}

SCREAMING_SNAKE_CASE							:Optional[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 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





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






	def __init__(       self					: Optional[int]							,A					: Optional[Any]							,A					: Optional[int]=False							,A					: int=False							,A					: Union[str, Any]=False							,A					: int=None							,A					: Optional[Any]=None							,A					: Union[str, Any]=None							,A					: Optional[Any]=None							,A					: Optional[Dict[str, Any]] = None							,**A					: Tuple							,):
					__A   			= {} if sp_model_kwargs is None else sp_model_kwargs

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

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

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

					__A   			= do_lower_case
					__A   			= remove_space
					__A   			= keep_accents
					__A   			= vocab_file

					__A   			= spm.SentencePieceProcessor(**self.sp_model_kwargs )
					self.sp_model.Load(A )

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

					# Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
					__A   			= re.compile(
					    f'''[{''.join(map(A							,list(range(0							,9 ) ) + list(range(11							,32 ) ) + list(range(1_27							,1_60 ) ) + [1_60, 1_73, 82_03] ) )}]''' )






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






	def __setstate__(       self					: Optional[Any]							,A					: Union[str, Any] ):
					__A   			= d

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

					__A   			= 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 				UpperCamelCase_  (       self					: List[str] ):
					return len(self.sp_model )






	def 				UpperCamelCase_  (       self					: int							,A					: str ):
					__A   			= self.non_printing_characters_re.sub(""							,A )

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

					# NFC Unicode normalization
					__A   			= unicodedata.normalize("NFC"							,A )
					return text






	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: str							,**A					: Optional[int] ):
					__A   			= self.preprocess_text(A )
					return self.sp_model.encode(A							,out_type=A )






	def 				UpperCamelCase_  (       self					: Any							,A					: str ):
					return self.sp_model.PieceToId(A )






	def 				UpperCamelCase_  (       self					: Dict							,A					: int ):
					return self.sp_model.IdToPiece(A )






	@staticmethod
	def 				UpperCamelCase_  (       A					: str ):
					return out_string






	def 				UpperCamelCase_  (       self					: str							,A					: List[str] ):
					__A   			= []
					__A   			= ""
					__A   			= 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(A ) + token
													__A   			= True
													__A   			= []
									else:
													current_sub_tokens.append(A )
													__A   			= False
					out_string += self.sp_model.decode(A )

					return out_string






	def 				UpperCamelCase_  (       self					: str ):
					__A   			= {self.convert_ids_to_tokens(A ): i for i in range(self.vocab_size )}
					vocab.update(self.added_tokens_encoder )
					return vocab






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

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

					return (out_vocab_file,)






	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: Union[str, List[str]]							,A					: Union[str, bool] = False ):

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

					if return_tensors is True or return_tensors == "pt":
									__A   			= torch.tensor(A )

					return token_ids






	def 				UpperCamelCase_  (       self					: List[Any]							,A					: Union[int, List[int]] ):
					return self.sp_model.decode(A )







	def 				UpperCamelCase_  (       self					: List[str]							,A					: "Conversation" ):
					__A   			= [f'''User: {text}''' if is_user else f'''Bot: {text}''' for is_user, text in conversation.iter_texts()]
					__A   			= (
					    f'''{self.eos_token}{self.bos_token}''' + f'''{self.bos_token}'''.join(A ) + f'''{self.bos_token}Bot:'''
					)
					return self.encode(text=A )
