


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


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

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

__snake_case   			=    {
    """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				(      snake_case_	):
   __UpperCAmelCase :       Dict							  =  VOCAB_FILES_NAMES
   __UpperCAmelCase :       int							  =  PRETRAINED_VOCAB_FILES_MAP
   __UpperCAmelCase :       int							  =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   __UpperCAmelCase :       int							  =  ['''input_ids''', '''attention_mask''']







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

       '''simple docstring'''




       snake_case    :      List[str]      					=    {} if sp_model_kwargs is None else sp_model_kwargs

       snake_case    :      Tuple      					=    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"						)
           snake_case    :      Optional[Any]      					=    "None"

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

       snake_case    :      Any      					=    do_lower_case
       snake_case    :      Tuple      					=    remove_space
       snake_case    :      Any      					=    keep_accents
       snake_case    :      Tuple      					=    vocab_file

       snake_case    :      Union[str, Any]      					=    spm.SentencePieceProcessor(**self.sp_model_kwargs						)
       self.sp_model.Load(UpperCamelCase__						)

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

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







   def __getstate__(	self						)	->		int:

       '''simple docstring'''




       snake_case    :      Optional[int]      					=    self.__dict__.copy()
       snake_case    :      List[str]      					=    None
       return state







   def __setstate__(	self						,	UpperCamelCase__						)	->		int:

       '''simple docstring'''




       snake_case    :      Dict      					=    d

       # for backward compatibility
       if not hasattr(self						,	"sp_model_kwargs"						):
           snake_case    :      List[str]      					=    {}

       snake_case    :      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  lowerCamelCase		(	self						)	->		int:

       '''simple docstring'''




       return len(self.sp_model						)







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		str:

       '''simple docstring'''




       snake_case    :      int      					=    self.non_printing_characters_re.sub(""						,	UpperCamelCase__						)

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

       # NFC Unicode normalization
       snake_case    :      Any      					=    unicodedata.normalize("NFC"						,	UpperCamelCase__						)
       return text







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	**UpperCamelCase__						)	->		List[str]:

       '''simple docstring'''




       snake_case    :      Dict      					=    self.preprocess_text(UpperCamelCase__						)
       return self.sp_model.encode(UpperCamelCase__						,	out_type=UpperCamelCase__						)







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		int:

       '''simple docstring'''




       return self.sp_model.PieceToId(UpperCamelCase__						)







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		str:

       '''simple docstring'''




       return self.sp_model.IdToPiece(UpperCamelCase__						)







   @staticmethod
   def  lowerCamelCase		(	UpperCamelCase__						)	->		str:

       '''simple docstring'''




       return out_string







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		str:

       '''simple docstring'''




       snake_case    :      Dict      					=    []
       snake_case    :      Any      					=    ""
       snake_case    :      List[str]      					=    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
               snake_case    :      List[str]      					=    True
               snake_case    :      Optional[int]      					=    []
           else:
               current_sub_tokens.append(UpperCamelCase__						)
               snake_case    :      List[str]      					=    False
       out_string += self.sp_model.decode(UpperCamelCase__						)

       return out_string







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

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    {self.convert_ids_to_tokens(UpperCamelCase__						): i for i in range(self.vocab_size						)}
       vocab.update(self.added_tokens_encoder						)
       return vocab







   def  lowerCamelCase		(	self						,	UpperCamelCase__						,	UpperCamelCase__ = None						)	->		Tuple[str]:

       '''simple docstring'''




       if not os.path.isdir(UpperCamelCase__						):
           logger.error(F'Vocabulary path ({save_directory}) should be a directory'						)
           return
       snake_case    :      Tuple      					=    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:
               snake_case    :      str      					=    self.sp_model.serialized_model_proto()
               fi.write(UpperCamelCase__						)

       return (out_vocab_file,)







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

       '''simple docstring'''




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

       if return_tensors is True or return_tensors == "pt":
           snake_case    :      Optional[Any]      					=    torch.tensor(UpperCamelCase__						)

       return token_ids







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		str:

       '''simple docstring'''




       return self.sp_model.decode(UpperCamelCase__						)







   def  lowerCamelCase		(	self						,	UpperCamelCase__						)	->		List[int]:

       '''simple docstring'''




       snake_case    :      int      					=    [F'User: {text}' if is_user else F'Bot: {text}' for is_user, text in conversation.iter_texts()]
       snake_case    :      Union[str, Any]      					=    (
           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 random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
  import onnxruntime as ort







class _lowerCAmelCase				(      snake_case_    ,							unittest.TestCase	):
   # TODO: is there an appropriate internal test set?
   __UpperCAmelCase :       Optional[Any]							  =  '''ssube/stable-diffusion-x4-upscaler-onnx'''







   def  lowerCamelCase		(	self						,	UpperCamelCase__=0						)	->		Tuple:

       '''simple docstring'''




       snake_case    :      Union[str, Any]      					=    floats_tensor((1, 3, 128, 128)						,	rng=random.Random(UpperCamelCase__						)						)
       snake_case    :      Any      					=    torch.manual_seed(UpperCamelCase__						)
       snake_case    :      List[str]      					=    {
           "prompt": "A painting of a squirrel eating a burger",
           "image": image,
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 7.5,
           "output_type": "numpy",
       }
       return inputs







   def  lowerCamelCase		(	self						)	->		Tuple:

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Tuple      					=    self.get_dummy_inputs()
       snake_case    :      int      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      Tuple      					=    image[0, -3:, -3:, -1].flatten()

       # started as 128, should now be 512
       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Optional[int]      					=    np.array(
           [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
       assert np.abs(image_slice - expected_slice						).max() < 1e-1







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

       '''simple docstring'''




       snake_case    :      Tuple      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Optional[Any]      					=    PNDMScheduler.from_config(pipe.scheduler.config						,	skip_prk_steps=UpperCamelCase__						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Tuple      					=    self.get_dummy_inputs()
       snake_case    :      Optional[Any]      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      List[str]      					=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Dict      					=    np.array(
           [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964]						)
       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1







   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       snake_case    :      Optional[int]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Optional[Any]      					=    DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Any      					=    self.get_dummy_inputs()
       snake_case    :      List[Any]      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      List[str]      					=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Optional[Any]      					=    np.array(
           [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1







   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




       snake_case    :      Union[str, Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Any      					=    EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Dict      					=    self.get_dummy_inputs()
       snake_case    :      Union[str, Any]      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      str      					=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Union[str, Any]      					=    np.array(
           [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223]						)
       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1







   def  lowerCamelCase		(	self						)	->		str:

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Optional[Any]      					=    EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Any      					=    self.get_dummy_inputs()
       snake_case    :      int      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      List[Any]      					=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Union[str, Any]      					=    np.array(
           [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043]						)

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-1







@nightly
@require_onnxruntime
@require_torch_gpu
class _lowerCAmelCase				(      unittest.TestCase	):







   @property
   def  lowerCamelCase		(	self						)	->		Tuple:

       '''simple docstring'''




       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )







   @property
   def  lowerCamelCase		(	self						)	->		List[Any]:

       '''simple docstring'''




       snake_case    :      List[Any]      					=    ort.SessionOptions()
       snake_case    :      str      					=    False
       return options







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

       '''simple docstring'''




       snake_case    :      str      					=    load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"						)
       snake_case    :      Tuple      					=    init_image.resize((128, 128)						)
       # using the PNDM scheduler by default
       snake_case    :      Tuple      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(
           "ssube/stable-diffusion-x4-upscaler-onnx"						,	provider=self.gpu_provider						,	sess_options=self.gpu_options						,	)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      List[Any]      					=    "A fantasy landscape, trending on artstation"

       snake_case    :      Union[str, Any]      					=    torch.manual_seed(0						)
       snake_case    :      int      					=    pipe(
           prompt=UpperCamelCase__						,	image=UpperCamelCase__						,	guidance_scale=7.5						,	num_inference_steps=10						,	generator=UpperCamelCase__						,	output_type="np"						,	)
       snake_case    :      str      					=    output.images
       snake_case    :      Tuple      					=    images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 512, 3)
       snake_case    :      int      					=    np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972]						)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2







   def  lowerCamelCase		(	self						)	->		Optional[Any]:

       '''simple docstring'''




       snake_case    :      Tuple      					=    load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"						)
       snake_case    :      Optional[Any]      					=    init_image.resize((128, 128)						)
       snake_case    :      Union[str, Any]      					=    LMSDiscreteScheduler.from_pretrained(
           "ssube/stable-diffusion-x4-upscaler-onnx"						,	subfolder="scheduler"						)
       snake_case    :      List[Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(
           "ssube/stable-diffusion-x4-upscaler-onnx"						,	scheduler=UpperCamelCase__						,	provider=self.gpu_provider						,	sess_options=self.gpu_options						,	)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Union[str, Any]      					=    "A fantasy landscape, trending on artstation"

       snake_case    :      Tuple      					=    torch.manual_seed(0						)
       snake_case    :      str      					=    pipe(
           prompt=UpperCamelCase__						,	image=UpperCamelCase__						,	guidance_scale=7.5						,	num_inference_steps=20						,	generator=UpperCamelCase__						,	output_type="np"						,	)
       snake_case    :      List[Any]      					=    output.images
       snake_case    :      List[Any]      					=    images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 512, 3)
       snake_case    :      int      					=    np.array(
           [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566]						)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

       assert np.abs(image_slice.flatten() - expected_slice						).max() < 2e-2





