





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


UpperCAmelCase     :     int							=	logging.get_logger(__name__)
UpperCAmelCase     :     List[Any]							=	{"""vocab_file""": """spiece.model"""}

UpperCAmelCase     :     int							=	{
    """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""",
    }
}

UpperCAmelCase     :     str							=	{
    """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	(   UpperCamelCase__):
      _lowercase    : int            =      VOCAB_FILES_NAMES
      _lowercase    : Optional[int]            =      PRETRAINED_VOCAB_FILES_MAP
      _lowercase    : Any            =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _lowercase    : List[str]            =      ["""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__		,				)       ->						None:



        '''simple docstring'''




        a__     : str  ={} if sp_model_kwargs is None else sp_model_kwargs

        a__     : Any  =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__     : int  ="None"

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

        a__     : List[str]  =do_lower_case
        a__     : Union[str, Any]  =remove_space
        a__     : Union[str, Any]  =keep_accents
        a__     : int  =vocab_file

        a__     : Any  =spm.SentencePieceProcessor(**self.sp_model_kwargs							)
        self.sp_model.Load(lowerCAmelCase__							)

        # Used for whitespace normalization in input texts
        # fmt : off
        a__     : List[Any]  ={" ", " ", " ", " ", " ", "　", " ", " ", " ", " ", "￼", ""}
        # fmt : on

        # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
        a__     : Optional[Any]  =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							)       ->						str:



        '''simple docstring'''




        a__     : Dict  =self.__dict__.copy()
        a__     : Optional[int]  =None
        return state







      def __setstate__(   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : List[Any]  =d

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

        a__     : 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        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        return len(self.sp_model							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        a__     : List[str]  =self.non_printing_characters_re.sub(""		,				lowerCAmelCase__							)

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

        # NFC Unicode normalization
        a__     : Optional[Any]  =unicodedata.normalize("NFC"		,				lowerCAmelCase__							)
        return text







      def        _lowercase  (   self		,				lowerCAmelCase__		,				**lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        a__     : Dict  =self.preprocess_text(lowerCAmelCase__							)
        return self.sp_model.encode(lowerCAmelCase__		,				out_type=lowerCAmelCase__							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        return self.sp_model.PieceToId(lowerCAmelCase__							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        return self.sp_model.IdToPiece(lowerCAmelCase__							)







      @staticmethod
      def        _lowercase  (   lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        return out_string







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




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

        return out_string







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



        '''simple docstring'''




        a__     : List[str]  ={self.convert_ids_to_tokens(lowerCAmelCase__							): i for i in range(self.vocab_size							)}
        vocab.update(self.added_tokens_encoder							)
        return vocab







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__ = None							)       ->						Tuple[str]:



        '''simple docstring'''




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

        return (out_vocab_file,)







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



        '''simple docstring'''




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

        if return_tensors is True or return_tensors == "pt":
          a__     : Optional[int]  =torch.tensor(lowerCAmelCase__							)

        return token_ids







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        return self.sp_model.decode(lowerCAmelCase__							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						List[int]:



        '''simple docstring'''




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













import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest







class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : Any            =      (PNDMScheduler,)
      _lowercase    : str            =      (("""num_inference_steps""", 50),)







      def        _lowercase  (   self		,				**lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : Dict  ={
            "num_train_timesteps": 1_0_0_0,
            "beta_start": 0.00_01,
            "beta_end": 0.02,
            "beta_schedule": "linear",
        }

        config.update(**lowerCAmelCase__							)
        return config







      def        _lowercase  (   self		,				lowerCAmelCase__=0		,				**lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        a__     : Optional[int]  =dict(self.forward_default_kwargs							)
        a__     : Tuple  =kwargs.pop("num_inference_steps"		,				lowerCAmelCase__							)
        a__     : List[str]  =self.dummy_sample
        a__     : List[str]  =0.1 * sample
        a__     : str  =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

        for scheduler_class in self.scheduler_classes:
          a__     : int  =self.get_scheduler_config(**lowerCAmelCase__							)
          a__     : Union[str, Any]  =scheduler_class(**lowerCAmelCase__							)
          scheduler.set_timesteps(lowerCAmelCase__							)
          # copy over dummy past residuals
          a__     : Any  =dummy_past_residuals[:]

          with tempfile.TemporaryDirectory() as tmpdirname:
            scheduler.save_config(lowerCAmelCase__							)
            a__     : List[Any]  =scheduler_class.from_pretrained(lowerCAmelCase__							)
            new_scheduler.set_timesteps(lowerCAmelCase__							)
            # copy over dummy past residuals
            a__     : str  =dummy_past_residuals[:]

          a__     : Any  =scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : Dict  =new_scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          assert torch.sum(torch.abs(output - new_output							)							) < 1E-5, "Scheduler outputs are not identical"

          a__     : Optional[int]  =scheduler.step_plms(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : Union[str, Any]  =new_scheduler.step_plms(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          assert torch.sum(torch.abs(output - new_output							)							) < 1E-5, "Scheduler outputs are not identical"







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        pass







      def        _lowercase  (   self		,				lowerCAmelCase__=0		,				**lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Optional[int]  =dict(self.forward_default_kwargs							)
        a__     : List[str]  =kwargs.pop("num_inference_steps"		,				lowerCAmelCase__							)
        a__     : List[str]  =self.dummy_sample
        a__     : int  =0.1 * sample
        a__     : Tuple  =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

        for scheduler_class in self.scheduler_classes:
          a__     : Dict  =self.get_scheduler_config()
          a__     : List[str]  =scheduler_class(**lowerCAmelCase__							)
          scheduler.set_timesteps(lowerCAmelCase__							)

          # copy over dummy past residuals (must be after setting timesteps)
          a__     : Dict  =dummy_past_residuals[:]

          with tempfile.TemporaryDirectory() as tmpdirname:
            scheduler.save_config(lowerCAmelCase__							)
            a__     : Dict  =scheduler_class.from_pretrained(lowerCAmelCase__							)
            # copy over dummy past residuals
            new_scheduler.set_timesteps(lowerCAmelCase__							)

            # copy over dummy past residual (must be after setting timesteps)
            a__     : Optional[int]  =dummy_past_residuals[:]

          a__     : Optional[Any]  =scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : List[Any]  =new_scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          assert torch.sum(torch.abs(output - new_output							)							) < 1E-5, "Scheduler outputs are not identical"

          a__     : List[str]  =scheduler.step_plms(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : Any  =new_scheduler.step_plms(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          assert torch.sum(torch.abs(output - new_output							)							) < 1E-5, "Scheduler outputs are not identical"







      def        _lowercase  (   self		,				**lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : Union[str, Any]  =self.scheduler_classes[0]
        a__     : Optional[Any]  =self.get_scheduler_config(**lowerCAmelCase__							)
        a__     : Any  =scheduler_class(**lowerCAmelCase__							)

        a__     : int  =1_0
        a__     : Union[str, Any]  =self.dummy_model()
        a__     : Optional[int]  =self.dummy_sample_deter
        scheduler.set_timesteps(lowerCAmelCase__							)

        for i, t in enumerate(scheduler.prk_timesteps							):
          a__     : List[Any]  =model(lowerCAmelCase__		,				lowerCAmelCase__							)
          a__     : Optional[Any]  =scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							).prev_sample

        for i, t in enumerate(scheduler.plms_timesteps							):
          a__     : int  =model(lowerCAmelCase__		,				lowerCAmelCase__							)
          a__     : int  =scheduler.step_plms(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							).prev_sample

        return sample







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



        '''simple docstring'''




        a__     : str  =dict(self.forward_default_kwargs							)

        a__     : Tuple  =kwargs.pop("num_inference_steps"		,				lowerCAmelCase__							)

        for scheduler_class in self.scheduler_classes:
          a__     : Union[str, Any]  =self.get_scheduler_config()
          a__     : List[str]  =scheduler_class(**lowerCAmelCase__							)

          a__     : List[Any]  =self.dummy_sample
          a__     : Dict  =0.1 * sample

          if num_inference_steps is not None and hasattr(lowerCAmelCase__		,				"set_timesteps"							):
            scheduler.set_timesteps(lowerCAmelCase__							)
          elif num_inference_steps is not None and not hasattr(lowerCAmelCase__		,				"set_timesteps"							):
            a__     : int  =num_inference_steps

          # copy over dummy past residuals (must be done after set_timesteps)
          a__     : Tuple  =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
          a__     : str  =dummy_past_residuals[:]

          a__     : List[Any]  =scheduler.step_prk(lowerCAmelCase__		,				0		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : int  =scheduler.step_prk(lowerCAmelCase__		,				1		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          self.assertEqual(output_a.shape		,				sample.shape							)
          self.assertEqual(output_a.shape		,				output_a.shape							)

          a__     : List[str]  =scheduler.step_plms(lowerCAmelCase__		,				0		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample
          a__     : Dict  =scheduler.step_plms(lowerCAmelCase__		,				1		,				lowerCAmelCase__		,				**lowerCAmelCase__							).prev_sample

          self.assertEqual(output_a.shape		,				sample.shape							)
          self.assertEqual(output_a.shape		,				output_a.shape							)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        for timesteps in [1_0_0, 1_0_0_0]:
          self.check_over_configs(num_train_timesteps=lowerCAmelCase__							)







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



        '''simple docstring'''




        for steps_offset in [0, 1]:
          self.check_over_configs(steps_offset=lowerCAmelCase__							)

        a__     : Optional[Any]  =self.scheduler_classes[0]
        a__     : Tuple  =self.get_scheduler_config(steps_offset=1							)
        a__     : Optional[Any]  =scheduler_class(**lowerCAmelCase__							)
        scheduler.set_timesteps(1_0							)
        assert torch.equal(
            scheduler.timesteps		,				torch.LongTensor(
                [9_0_1, 8_5_1, 8_5_1, 8_0_1, 8_0_1, 7_5_1, 7_5_1, 7_0_1, 7_0_1, 6_5_1, 6_5_1, 6_0_1, 6_0_1, 5_0_1, 4_0_1, 3_0_1, 2_0_1, 1_0_1, 1]							)		,				)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        for beta_start, beta_end in zip([0.00_01, 0.0_01]		,				[0.0_02, 0.02]							):
          self.check_over_configs(beta_start=lowerCAmelCase__		,				beta_end=lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        for schedule in ["linear", "squaredcos_cap_v2"]:
          self.check_over_configs(beta_schedule=lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        for prediction_type in ["epsilon", "v_prediction"]:
          self.check_over_configs(prediction_type=lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        for t in [1, 5, 1_0]:
          self.check_over_forward(time_step=lowerCAmelCase__							)







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



        '''simple docstring'''




        for t, num_inference_steps in zip([1, 5, 1_0]		,				[1_0, 5_0, 1_0_0]							):
          self.check_over_forward(num_inference_steps=lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : Dict  =2_7

        for scheduler_class in self.scheduler_classes:
          a__     : Tuple  =self.dummy_sample
          a__     : Dict  =0.1 * sample

          a__     : Dict  =self.get_scheduler_config()
          a__     : int  =scheduler_class(**lowerCAmelCase__							)

          scheduler.set_timesteps(lowerCAmelCase__							)

          # before power of 3 fix, would error on first step, so we only need to do two
          for i, t in enumerate(scheduler.prk_timesteps[:2]							):
            a__     : Any  =scheduler.step_prk(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							).prev_sample







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



        '''simple docstring'''




        with self.assertRaises(lowerCAmelCase__							):
          a__     : List[Any]  =self.scheduler_classes[0]
          a__     : Dict  =self.get_scheduler_config()
          a__     : Tuple  =scheduler_class(**lowerCAmelCase__							)

          scheduler.step_plms(self.dummy_sample		,				1		,				self.dummy_sample							).prev_sample







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



        '''simple docstring'''




        a__     : List[Any]  =self.full_loop()
        a__     : str  =torch.sum(torch.abs(lowerCAmelCase__							)							)
        a__     : Optional[Any]  =torch.mean(torch.abs(lowerCAmelCase__							)							)

        assert abs(result_sum.item() - 1_98.13_18							) < 1E-2
        assert abs(result_mean.item() - 0.25_80							) < 1E-3







      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : str  =self.full_loop(prediction_type="v_prediction"							)
        a__     : int  =torch.sum(torch.abs(lowerCAmelCase__							)							)
        a__     : Optional[int]  =torch.mean(torch.abs(lowerCAmelCase__							)							)

        assert abs(result_sum.item() - 67.39_86							) < 1E-2
        assert abs(result_mean.item() - 0.08_78							) < 1E-3







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Tuple  =self.full_loop(set_alpha_to_one=lowerCAmelCase__		,				beta_start=0.01							)
        a__     : str  =torch.sum(torch.abs(lowerCAmelCase__							)							)
        a__     : Dict  =torch.mean(torch.abs(lowerCAmelCase__							)							)

        assert abs(result_sum.item() - 2_30.03_99							) < 1E-2
        assert abs(result_mean.item() - 0.29_95							) < 1E-3







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



        '''simple docstring'''




        a__     : Dict  =self.full_loop(set_alpha_to_one=lowerCAmelCase__		,				beta_start=0.01							)
        a__     : Union[str, Any]  =torch.sum(torch.abs(lowerCAmelCase__							)							)
        a__     : Union[str, Any]  =torch.mean(torch.abs(lowerCAmelCase__							)							)

        assert abs(result_sum.item() - 1_86.94_82							) < 1E-2
        assert abs(result_mean.item() - 0.24_34							) < 1E-3






