
'''simple docstring'''


import importlib
import math
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Tuple, Union

import flax
import jax.numpy as jnp

from ..utils import BaseOutput


snake_case_ :  Tuple     =					'scheduler_config.json'


class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							1
   _snake_case		   =							2
   _snake_case		   =							3
   _snake_case		   =							4
   _snake_case		   =							5




@dataclass
class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							42




class 						lowercase__	:






   '''simple docstring'''







   _snake_case		   =							SCHEDULER_CONFIG_NAME
   _snake_case		   =							['''dtype''']
   _snake_case		   =							[]
   _snake_case		   =							True







   @classmethod
   def       UpperCAmelCase				(			cls ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__=False ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  cls.load_config(
          pretrained_model_name_or_path=lowerCamelCase__ ,   subfolder=lowerCamelCase__ ,   return_unused_kwargs=lowerCamelCase__ ,   **lowerCamelCase__ ,   )
      UpperCamelCase     , UpperCamelCase       =  cls.from_config(lowerCamelCase__ ,   return_unused_kwargs=lowerCamelCase__ ,   **lowerCamelCase__   )

      if hasattr(lowerCamelCase__ ,   '''create_state'''   ) and getattr(lowerCamelCase__ ,   '''has_state''' ,   lowerCamelCase__   ):
         UpperCamelCase       =  scheduler.create_state()

      if return_unused_kwargs:
         return scheduler, state, unused_kwargs

      return scheduler, state







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      self.save_config(save_directory=lowerCamelCase__ ,   push_to_hub=lowerCamelCase__ ,   **lowerCamelCase__   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return self._get_compatibles()







   @classmethod
   def       UpperCAmelCase				(			cls   ):

      '''simple docstring'''







      UpperCamelCase       =  list(set([cls.__name__] + cls._compatibles   )   )
      UpperCamelCase       =  importlib.import_module(__name__.split('''.'''   )[0]   )
      UpperCamelCase       =  [
          getattr(lowerCamelCase__ ,   lowerCamelCase__   ) for c in compatible_classes_str if hasattr(lowerCamelCase__ ,   lowerCamelCase__   )
      ]
      return compatible_classes




def  __snake_case     (    _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       Tuple[int]):
   assert len(_UpperCAmelCase) >= x.ndim
   return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(_UpperCAmelCase) - x.ndim)),      _UpperCAmelCase)



def  __snake_case     (    _UpperCAmelCase					:       int,      _UpperCAmelCase					:       Optional[int]=0.9_9_9,      _UpperCAmelCase					:       Tuple=jnp.floataa):

   def alpha_bar(_UpperCAmelCase					:       Optional[Any]):
      return math.cos((time_step + 0.0_0_8) / 1.0_0_8 * math.pi / 2) ** 2

   UpperCamelCase       =  []
   for i in range(_UpperCAmelCase):
      UpperCamelCase       =  i / num_diffusion_timesteps
      UpperCamelCase       =  (i + 1) / num_diffusion_timesteps
      betas.append(min(1 - alpha_bar(_UpperCAmelCase) / alpha_bar(_UpperCAmelCase),      _UpperCAmelCase))
   return jnp.array(_UpperCAmelCase,      dtype=_UpperCAmelCase)


@flax.struct.dataclass
class 						lowercase__	:






   '''simple docstring'''







   _snake_case		   =							42
   _snake_case		   =							42
   _snake_case		   =							42







   @classmethod
   def       UpperCAmelCase				(			cls ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  scheduler.config

      if config.trained_betas is not None:
         UpperCamelCase       =  jnp.asarray(config.trained_betas ,   dtype=scheduler.dtype   )
      elif config.beta_schedule == "linear":
         UpperCamelCase       =  jnp.linspace(config.beta_start ,   config.beta_end ,   config.num_train_timesteps ,   dtype=scheduler.dtype   )
      elif config.beta_schedule == "scaled_linear":
         # this schedule is very specific to the latent diffusion model.
         UpperCamelCase       =  (
             jnp.linspace(
                 config.beta_start**0.5 ,   config.beta_end**0.5 ,   config.num_train_timesteps ,   dtype=scheduler.dtype   )
             ** 2
         )
      elif config.beta_schedule == "squaredcos_cap_v2":
         # Glide cosine schedule
         UpperCamelCase       =  betas_for_alpha_bar(config.num_train_timesteps ,   dtype=scheduler.dtype   )
      else:
         raise NotImplementedError(
             f'beta_schedule {config.beta_schedule} is not implemented for scheduler {scheduler.__class__.__name__}'   )

      UpperCamelCase       =  1.0 - betas

      UpperCamelCase       =  jnp.cumprod(lowerCamelCase__ ,   axis=0   )

      return cls(
          alphas=lowerCamelCase__ ,   betas=lowerCamelCase__ ,   alphas_cumprod=lowerCamelCase__ ,   )




def  __snake_case     (    _UpperCAmelCase					:       CommonSchedulerState,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray):
   UpperCamelCase       =  state.alphas_cumprod

   UpperCamelCase       =  alphas_cumprod[timesteps] ** 0.5
   UpperCamelCase       =  sqrt_alpha_prod.flatten()
   UpperCamelCase       =  broadcast_to_shape_from_left(_UpperCAmelCase,      original_samples.shape)

   UpperCamelCase       =  (1 - alphas_cumprod[timesteps]) ** 0.5
   UpperCamelCase       =  sqrt_one_minus_alpha_prod.flatten()
   UpperCamelCase       =  broadcast_to_shape_from_left(_UpperCAmelCase,      original_samples.shape)

   return sqrt_alpha_prod, sqrt_one_minus_alpha_prod



def  __snake_case     (    _UpperCAmelCase					:       CommonSchedulerState,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray):
   UpperCamelCase     , UpperCamelCase       =  get_sqrt_alpha_prod(_UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase)
   UpperCamelCase       =  sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
   return noisy_samples



def  __snake_case     (    _UpperCAmelCase					:       CommonSchedulerState,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray,      _UpperCAmelCase					:       jnp.ndarray):
   UpperCamelCase     , UpperCamelCase       =  get_sqrt_alpha_prod(_UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase)
   UpperCamelCase       =  sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
   return velocity







'''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_ :  Union[str, Any]     =					logging.get_logger(__name__)
snake_case_ :  Union[str, Any]     =					{'vocab_file': 'spiece.model'}

snake_case_ :  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',
    }
}

snake_case_ :  Dict     =					{
    '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 						lowercase__	(  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 ,   lowerCamelCase__ ,   lowerCamelCase__=False ,   lowerCamelCase__=False ,   lowerCamelCase__=False ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  {} if sp_model_kwargs is None else sp_model_kwargs

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

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

      UpperCamelCase       =  do_lower_case
      UpperCamelCase       =  remove_space
      UpperCamelCase       =  keep_accents
      UpperCamelCase       =  vocab_file

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

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

      # Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
      UpperCamelCase       =  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'''







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







   def __setstate__(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  d

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

      UpperCamelCase       =  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   ):

      '''simple docstring'''







      return len(self.sp_model   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







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

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

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   **lowerCamelCase__   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.PieceToId(lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.IdToPiece(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__   ):

      '''simple docstring'''







      return out_string







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







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

      return out_string







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = None   ):

      '''simple docstring'''







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

      return (out_vocab_file,)







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = False   ):

      '''simple docstring'''







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

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

      return token_ids







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return self.sp_model.decode(lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







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





