





'''simple docstring'''

import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
    from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
    from tensorflow.keras.optimizers import Adam







class __a			( tf.keras.optimizers.schedules.LearningRateSchedule      ):





   def __init__(       self					:      List[str]					,lowerCamelCase					:      float					,lowerCamelCase					:      Callable					,lowerCamelCase					:      int					,lowerCamelCase					:      float = 1.0					,lowerCamelCase					:      str = None					,):

      '''simple docstring'''

      super().__init__()
      __SCREAMING_SNAKE_CASE   				=							initial_learning_rate
      __SCREAMING_SNAKE_CASE   				=							warmup_steps
      __SCREAMING_SNAKE_CASE   				=							power
      __SCREAMING_SNAKE_CASE   				=							decay_schedule_fn
      __SCREAMING_SNAKE_CASE   				=							name





   def __call__(       self					:      List[Any]					,lowerCamelCase					:      Optional[Any]     ):

      '''simple docstring'''

      with tf.name_scope(self.name or """WarmUp"""     ) as name:
         # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
         # learning rate will be `global_step/num_warmup_steps * init_lr`.
         __SCREAMING_SNAKE_CASE   				=							tf.cast(lowerCamelCase					,tf.floataa     )
         __SCREAMING_SNAKE_CASE   				=							tf.cast(self.warmup_steps					,tf.floataa     )
         __SCREAMING_SNAKE_CASE   				=							global_step_float / warmup_steps_float
         __SCREAMING_SNAKE_CASE   				=							self.initial_learning_rate * tf.math.pow(lowerCamelCase					,self.power     )
         return tf.cond(
             global_step_float < warmup_steps_float					,lambda: warmup_learning_rate					,lambda: self.decay_schedule_fn(step - self.warmup_steps     )					,name=lowerCamelCase					,)





   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      return {
          "initial_learning_rate": self.initial_learning_rate,
          "decay_schedule_fn": self.decay_schedule_fn,
          "warmup_steps": self.warmup_steps,
          "power": self.power,
          "name": self.name,
      }





def     __magic_name__							(				__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase				,						__UpperCAmelCase = 0.0				,						__UpperCAmelCase = 0.9				,						__UpperCAmelCase = 0.9_9_9				,						__UpperCAmelCase = 1e-8				,						__UpperCAmelCase = None				,						__UpperCAmelCase = None				,						__UpperCAmelCase = 0.0				,						__UpperCAmelCase = 1.0				,						__UpperCAmelCase = None				,						)    ->       Any:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE   				=							tf.keras.optimizers.schedules.PolynomialDecay(
       initial_learning_rate=__UpperCAmelCase				,						decay_steps=num_train_steps - num_warmup_steps				,						end_learning_rate=init_lr * min_lr_ratio				,						power=__UpperCAmelCase				,						)
   if num_warmup_steps:
      __SCREAMING_SNAKE_CASE   				=							WarmUp(
          initial_learning_rate=__UpperCAmelCase				,						decay_schedule_fn=__UpperCAmelCase				,						warmup_steps=__UpperCAmelCase				,						)
   if weight_decay_rate > 0.0:
      __SCREAMING_SNAKE_CASE   				=							AdamWeightDecay(
          learning_rate=__UpperCAmelCase				,						weight_decay_rate=__UpperCAmelCase				,						beta_a=__UpperCAmelCase				,						beta_a=__UpperCAmelCase				,						epsilon=__UpperCAmelCase				,						clipnorm=__UpperCAmelCase				,						global_clipnorm=__UpperCAmelCase				,						exclude_from_weight_decay=["""LayerNorm""", """layer_norm""", """bias"""]				,						include_in_weight_decay=__UpperCAmelCase				,						)
   else:
      __SCREAMING_SNAKE_CASE   				=							tf.keras.optimizers.Adam(
          learning_rate=__UpperCAmelCase				,						beta_a=__UpperCAmelCase				,						beta_a=__UpperCAmelCase				,						epsilon=__UpperCAmelCase				,						clipnorm=__UpperCAmelCase				,						global_clipnorm=__UpperCAmelCase				,						)
   # We return the optimizer and the LR scheduler in order to better track the
   # evolution of the LR independently of the optimizer.
   return optimizer, lr_schedule







class __a			( _snake_case      ):





   def __init__(       self					:      Dict					,lowerCamelCase					:      Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001					,lowerCamelCase					:      float = 0.9					,lowerCamelCase					:      float = 0.999					,lowerCamelCase					:      float = 1E-7					,lowerCamelCase					:      bool = False					,lowerCamelCase					:      float = 0.0					,lowerCamelCase					:      Optional[List[str]] = None					,lowerCamelCase					:      Optional[List[str]] = None					,lowerCamelCase					:      str = "AdamWeightDecay"					,**lowerCamelCase					:      Any					,):

      '''simple docstring'''

      super().__init__(lowerCamelCase					,lowerCamelCase					,lowerCamelCase					,lowerCamelCase					,lowerCamelCase					,lowerCamelCase					,**lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							weight_decay_rate
      __SCREAMING_SNAKE_CASE   				=							include_in_weight_decay
      __SCREAMING_SNAKE_CASE   				=							exclude_from_weight_decay





   @classmethod
   def   UpperCAmelCase__  (       cls					:      List[Any]					,lowerCamelCase					:      Optional[Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							{"""WarmUp""": WarmUp}
      return super(lowerCamelCase					,cls     ).from_config(lowerCamelCase					,custom_objects=lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Any					,lowerCamelCase					:      Union[str, Any]					,lowerCamelCase					:      Union[str, Any]					,lowerCamelCase					:      List[Any]     ):

      '''simple docstring'''

      super(lowerCamelCase					,self     )._prepare_local(lowerCamelCase					,lowerCamelCase					,lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							tf.constant(
          self.weight_decay_rate					,name="""adam_weight_decay_rate"""     )





   def   UpperCAmelCase__  (       self					:      List[str]					,lowerCamelCase					:      Union[str, Any]					,lowerCamelCase					:      Dict					,lowerCamelCase					:      Union[str, Any]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							self._do_use_weight_decay(var.name     )
      if do_decay:
         return var.assign_sub(
             learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["""weight_decay_rate"""]					,use_locking=self._use_locking					,)
      return tf.no_op()





   def   UpperCAmelCase__  (       self					:      Optional[int]					,lowerCamelCase					:      List[Any]					,lowerCamelCase					:      int=None					,**lowerCamelCase					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							list(zip(*lowerCamelCase     )     )
      return super(lowerCamelCase					,self     ).apply_gradients(zip(lowerCamelCase					,lowerCamelCase     )					,name=lowerCamelCase					,**lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Optional[Any]					,lowerCamelCase					:      Dict					,lowerCamelCase					:      Union[str, Any]					,lowerCamelCase					:      Any     ):

      '''simple docstring'''

      if apply_state is None:
         return self._decayed_lr_t[var_dtype], {}

      __SCREAMING_SNAKE_CASE   				=							apply_state or {}
      __SCREAMING_SNAKE_CASE   				=							apply_state.get((var_device, var_dtype)     )
      if coefficients is None:
         __SCREAMING_SNAKE_CASE   				=							self._fallback_apply_state(lowerCamelCase					,lowerCamelCase     )
         __SCREAMING_SNAKE_CASE   				=							coefficients

      return coefficients["lr_t"], {"apply_state": apply_state}





   def   UpperCAmelCase__  (       self					:      Optional[int]					,lowerCamelCase					:      Optional[Any]					,lowerCamelCase					:      int					,lowerCamelCase					:      Union[str, Any]=None     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							self._get_lr(var.device					,var.dtype.base_dtype					,lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							self._decay_weights_op(lowerCamelCase					,lowerCamelCase					,lowerCamelCase     )
      with tf.control_dependencies([decay]     ):
         return super(lowerCamelCase					,self     )._resource_apply_dense(lowerCamelCase					,lowerCamelCase					,**lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Union[str, Any]					,lowerCamelCase					:      Dict					,lowerCamelCase					:      Optional[Any]					,lowerCamelCase					:      List[str]					,lowerCamelCase					:      Tuple=None     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE   				=							self._get_lr(var.device					,var.dtype.base_dtype					,lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							self._decay_weights_op(lowerCamelCase					,lowerCamelCase					,lowerCamelCase     )
      with tf.control_dependencies([decay]     ):
         return super(lowerCamelCase					,self     )._resource_apply_sparse(lowerCamelCase					,lowerCamelCase					,lowerCamelCase					,**lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							super().get_config()
      config.update({"""weight_decay_rate""": self.weight_decay_rate}     )
      return config





   def   UpperCAmelCase__  (       self					:      int					,lowerCamelCase					:      List[Any]     ):

      '''simple docstring'''

      if self.weight_decay_rate == 0:
         return False

      if self._include_in_weight_decay:
         for r in self._include_in_weight_decay:
            if re.search(lowerCamelCase					,lowerCamelCase     ) is not None:
               return True

      if self._exclude_from_weight_decay:
         for r in self._exclude_from_weight_decay:
            if re.search(lowerCamelCase					,lowerCamelCase     ) is not None:
               return False
      return True







class __a			( _snake_case      ):





   def __init__(       self					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							[]
      __SCREAMING_SNAKE_CASE   				=							None





   @property
   def   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      if self._accum_steps is None:
         __SCREAMING_SNAKE_CASE   				=							tf.Variable(
             tf.constant(0					,dtype=tf.intaa     )					,trainable=lowerCamelCase					,synchronization=tf.VariableSynchronization.ON_READ					,aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA					,)

      return self._accum_steps.value()





   @property
   def   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

      if not self._gradients:
         raise ValueError("""The accumulator should be called first to initialize the gradients"""     )
      return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]





   def __call__(       self					:      Union[str, Any]					,lowerCamelCase					:      List[str]     ):

      '''simple docstring'''

      if not self._gradients:
         __SCREAMING_SNAKE_CASE   				=							self.step  # Create the step variable.
         self._gradients.extend(
             [
                 tf.Variable(
                     tf.zeros_like(lowerCamelCase     )					,trainable=lowerCamelCase					,synchronization=tf.VariableSynchronization.ON_READ					,aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA					,)
                 if gradient is not None
                 else gradient
                 for gradient in gradients
             ]     )
      if len(lowerCamelCase     ) != len(self._gradients     ):
         raise ValueError(f"""Expected {len(self._gradients     )} gradients, but got {len(lowerCamelCase     )}"""     )

      for accum_gradient, gradient in zip(self._gradients					,lowerCamelCase     ):
         if accum_gradient is not None and gradient is not None:
            accum_gradient.assign_add(lowerCamelCase     )

      self._accum_steps.assign_add(1     )





   def   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      if not self._gradients:
         return
      self._accum_steps.assign(0     )
      for gradient in self._gradients:
         if gradient is not None:
            gradient.assign(tf.zeros_like(lowerCamelCase     )     )










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_								=       logging.get_logger(__name__)
a_								=       {"""vocab_file""": """spiece.model"""}

a_								=       {
    """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_								=       {
    """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 			__lowerCAmelCase  (						lowerCAmelCase__ ):
 lowerCAmelCase__							       =	VOCAB_FILES_NAMES
 lowerCAmelCase__							       =	PRETRAINED_VOCAB_FILES_MAP
 lowerCAmelCase__							       =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowerCAmelCase__							       =	["""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						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     {} if sp_model_kwargs is None else sp_model_kwargs

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

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

    __lowerCamelCase    				=     do_lower_case
    __lowerCamelCase    				=     remove_space
    __lowerCamelCase    				=     keep_accents
    __lowerCamelCase    				=     vocab_file

    __lowerCamelCase    				=     spm.SentencePieceProcessor(**self.sp_model_kwargs					)
    self.sp_model.Load(__UpperCAmelCase					)

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

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




 def __getstate__(    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.__dict__.copy()
    __lowerCamelCase    				=     None
    return state




 def __setstate__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     d

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

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



    '''simple docstring'''
    return len(self.sp_model					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.non_printing_characters_re.sub(''''''						,					__UpperCAmelCase					)

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

    # NFC Unicode normalization
    __lowerCamelCase    				=     unicodedata.normalize('''NFC'''						,					__UpperCAmelCase					)
    return text




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					**__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.preprocess_text(__UpperCAmelCase					)
    return self.sp_model.encode(__UpperCAmelCase						,					out_type=__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    return self.sp_model.PieceToId(__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    return self.sp_model.IdToPiece(__UpperCAmelCase					)




 @staticmethod
 def 							lowerCamelCase   (    __UpperCAmelCase					):



    '''simple docstring'''
    return out_string




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



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

    return out_string




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     {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					):



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

    return (out_vocab_file,)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase = False					):



    '''simple docstring'''

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

    if return_tensors is True or return_tensors == "pt":
       __lowerCamelCase    				=     torch.tensor(__UpperCAmelCase					)

    return token_ids




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    return self.sp_model.decode(__UpperCAmelCase					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase					):



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






