



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


_lowerCAmelCase      						=	logging.get_logger(__name__)
_lowerCAmelCase      						=	{'''vocab_file''': '''spiece.model'''}

_lowerCAmelCase      						=	{
    '''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''',
    }
}

_lowerCAmelCase      						=	{
    '''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_(				SCREAMING_SNAKE_CASE_     ):
  '''simple docstring'''




  __lowercase							:      Dict				=      VOCAB_FILES_NAMES
  __lowercase							:      str				=      PRETRAINED_VOCAB_FILES_MAP
  __lowercase							:      Dict				=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  __lowercase							:      Optional[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:
      lowerCAmelCase__ :  Optional[int]							      =					{} if sp_model_kwargs is None else sp_model_kwargs

      lowerCAmelCase__ :  Dict							      =					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__ :  Tuple							      =					"""None"""

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

      lowerCAmelCase__ :  str							      =					spm.SentencePieceProcessor(**self.sp_model_kwargs      )
      self.sp_model.Load(__UpperCAmelCase      )

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

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




  def __getstate__(    self      )    ->		Any:
      lowerCAmelCase__ :  int							      =					self.__dict__.copy()
      lowerCAmelCase__ :  Optional[int]							      =					None
      return state




  def __setstate__(    self  ,__UpperCAmelCase      )    ->		List[str]:
      lowerCAmelCase__ :  List[str]							      =					d

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

      lowerCAmelCase__ :  Optional[int]							      =					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      )    ->		int:
      return len(self.sp_model      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		str:
      lowerCAmelCase__ :  Tuple							      =					self.non_printing_characters_re.sub(""""""  ,__UpperCAmelCase      )

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

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




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,**__UpperCAmelCase      )    ->		List[str]:
      lowerCAmelCase__ :  List[Any]							      =					self.preprocess_text(__UpperCAmelCase      )
      return self.sp_model.encode(__UpperCAmelCase  ,out_type=__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		int:
      return self.sp_model.PieceToId(__UpperCAmelCase      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		str:
      return self.sp_model.IdToPiece(__UpperCAmelCase      )




  @staticmethod
  def 							UpperCAmelCase_	(    __UpperCAmelCase      )    ->		str:
      return out_string




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		str:
      lowerCAmelCase__ :  int							      =					[]
      lowerCAmelCase__ :  Optional[int]							      =					""""""
      lowerCAmelCase__ :  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(__UpperCAmelCase      ) + token
              lowerCAmelCase__ :  Union[str, Any]							      =					True
              lowerCAmelCase__ :  Optional[Any]							      =					[]
          else:
              current_sub_tokens.append(__UpperCAmelCase      )
              lowerCAmelCase__ :  Any							      =					False
      out_string += self.sp_model.decode(__UpperCAmelCase      )

      return out_string




  def 							UpperCAmelCase_	(    self      )    ->		Dict[str, int]:
      lowerCAmelCase__ :  Optional[int]							      =					{self.convert_ids_to_tokens(__UpperCAmelCase      ): i for i in range(self.vocab_size      )}
      vocab.update(self.added_tokens_encoder      )
      return vocab




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Tuple[str]:
      if not os.path.isdir(__UpperCAmelCase      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
      lowerCAmelCase__ :  Optional[int]							      =					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__ :  str							      =					self.sp_model.serialized_model_proto()
              fi.write(__UpperCAmelCase      )

      return (out_vocab_file,)




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

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

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

      return token_ids




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		str:
      return self.sp_model.decode(__UpperCAmelCase      )






  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[int]:
      lowerCAmelCase__ :  List[Any]							      =					[F"""User: {text}""" if is_user else F"""Bot: {text}""" for is_user, text in conversation.iter_texts()]
      lowerCAmelCase__ :  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'''





from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


_lowerCAmelCase      						=	logging.get_logger(__name__)
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    try:
        with open(UpperCamelCase							,    """rb"""			) as flax_state_f:
            lowerCAmelCase__ :  Union[str, Any]							      =					from_bytes(UpperCamelCase							,    flax_state_f.read()			)
    except UnpicklingError as e:
        try:
            with open(UpperCamelCase			) as f:
                if f.read().startswith("""version"""			):
                    raise OSError(
                        """You seem to have cloned a repository without having git-lfs installed. Please"""
                        """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
                        """ folder you cloned."""			)
                else:
                    raise ValueError from e
        except (UnicodeDecodeError, ValueError):
            raise EnvironmentError(f"""Unable to convert {model_file} to Flax deserializable object. """			)

    return load_flax_weights_in_pytorch_model(UpperCamelCase							,    UpperCamelCase			)




def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    try:
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
            """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
            """ instructions."""			)
        raise

    # check if we have bf16 weights
    lowerCAmelCase__ :  str							      =					flatten_dict(jax.tree_util.tree_map(lambda UpperCamelCase			: x.dtype == jnp.bfloataa							,    UpperCamelCase			)			).values()
    if any(UpperCamelCase			):
        # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

        # and bf16 is not fully supported in PT yet.
        logger.warning(
            """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
            """before loading those in PyTorch model."""			)
        lowerCAmelCase__ :  Dict							      =					jax.tree_util.tree_map(
            lambda UpperCamelCase			: params.astype(np.floataa			) if params.dtype == jnp.bfloataa else params							,    UpperCamelCase			)

    lowerCAmelCase__ :  Any							      =					""""""

    lowerCAmelCase__ :  Any							      =					flatten_dict(UpperCamelCase							,    sep="""."""			)
    lowerCAmelCase__ :  Optional[int]							      =					pt_model.state_dict()

    # keep track of unexpected & missing keys
    lowerCAmelCase__ :  Optional[Any]							      =					[]
    lowerCAmelCase__ :  int							      =					set(pt_model_dict.keys()			)

    for flax_key_tuple, flax_tensor in flax_state_dict.items():
        lowerCAmelCase__ :  Union[str, Any]							      =					flax_key_tuple.split("""."""			)

        if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
            lowerCAmelCase__ :  Optional[int]							      =					flax_key_tuple_array[:-1] + ["""weight"""]
            lowerCAmelCase__ :  Any							      =					jnp.transpose(UpperCamelCase							,    (3, 2, 0, 1)			)
        elif flax_key_tuple_array[-1] == "kernel":
            lowerCAmelCase__ :  str							      =					flax_key_tuple_array[:-1] + ["""weight"""]
            lowerCAmelCase__ :  Any							      =					flax_tensor.T
        elif flax_key_tuple_array[-1] == "scale":
            lowerCAmelCase__ :  int							      =					flax_key_tuple_array[:-1] + ["""weight"""]

        if "time_embedding" not in flax_key_tuple_array:
            for i, flax_key_tuple_string in enumerate(UpperCamelCase			):
                lowerCAmelCase__ :  List[str]							      =					(
                    flax_key_tuple_string.replace("""_0"""							,    """.0"""			)
                    .replace("""_1"""							,    """.1"""			)
                    .replace("""_2"""							,    """.2"""			)
                    .replace("""_3"""							,    """.3"""			)
                    .replace("""_4"""							,    """.4"""			)
                    .replace("""_5"""							,    """.5"""			)
                    .replace("""_6"""							,    """.6"""			)
                    .replace("""_7"""							,    """.7"""			)
                    .replace("""_8"""							,    """.8"""			)
                    .replace("""_9"""							,    """.9"""			)
                )

        lowerCAmelCase__ :  Union[str, Any]							      =					""".""".join(UpperCamelCase			)

        if flax_key in pt_model_dict:
            if flax_tensor.shape != pt_model_dict[flax_key].shape:
                raise ValueError(
                    f"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
                    f"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""			)
            else:
                # add weight to pytorch dict
                lowerCAmelCase__ :  int							      =					np.asarray(UpperCamelCase			) if not isinstance(UpperCamelCase							,    np.ndarray			) else flax_tensor
                lowerCAmelCase__ :  int							      =					torch.from_numpy(UpperCamelCase			)
                # remove from missing keys
                missing_keys.remove(UpperCamelCase			)
        else:
            # weight is not expected by PyTorch model
            unexpected_keys.append(UpperCamelCase			)

    pt_model.load_state_dict(UpperCamelCase			)

    # re-transform missing_keys to list
    lowerCAmelCase__ :  Optional[int]							      =					list(UpperCamelCase			)

    if len(UpperCamelCase			) > 0:
        logger.warning(
            """Some weights of the Flax model were not used when initializing the PyTorch model"""
            f""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
            f""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
            """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
            f""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
            """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
            """ FlaxBertForSequenceClassification model)."""			)
    if len(UpperCamelCase			) > 0:
        logger.warning(
            f"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
            f""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
            """ use it for predictions and inference."""			)

    return pt_model


