the tokenizer used by the gptsw3 models import os import re import unicodedata from shutil import copyfile from typing import any dict list optional tuple union import sentencepiece as spm from tokenizationutils import pretrainedtokenizer from utils import istorchavailable logging if istorchavailable import torch logger logging getloggername vocabfilesnames vocabfile spiece model pretrainedvocabfilesmap vocabfile aiswedengptsw3126m https huggingface coaiswedengptsw3126mresolvemainspiece model aiswedengptsw3350m https huggingface coaiswedengptsw3350mresolvemainspiece model aiswedengptsw31 6b https huggingface coaiswedengptsw31 6bresolvemainspiece model aiswedengptsw36 7b https huggingface coaiswedengptsw36 7bresolvemainspiece model aiswedengptsw320b https huggingface coaiswedengptsw320bresolvemainspiece model pretrainedpositionalembeddingssizes aiswedengptsw3126m 2048 aiswedengptsw3350m 2048 aiswedengptsw31 6b 2048 aiswedengptsw36 7b 2048 aiswedengptsw320b 2048 class gptsw3tokenizerpretrainedtokenizer vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxmodelinputsizes pretrainedpositionalembeddingssizes modelinputnames inputids tokentypeids attentionmask def init self vocabfile dolowercasefalse removespacefalse keepaccentsfalse padtokennone unktokennone eostokennone bostokennone spmodelkwargs optionaldictstr any none kwargs none self spmodelkwargs if spmodelkwargs is none else spmodelkwargs nameorpath kwargs getnameorpath if nameorpath is none logger warning nameorpath not provided will work for all gptsw3 models except gptsw37b you are testing the model this can safely be ignored nameorpath none default definitions for our 2 tokenizer versions with nonechecks to enable proper testing eostoken endoftext if eostoken is none else eostoken unktoken unk if unktoken is none else unktoken if gptsw37b in nameorpath padtoken unktoken if padtoken is none else padtoken bostoken eostoken if bostoken is none else bostoken else padtoken pad if padtoken is none else padtoken bostoken s if bostoken is none else bostoken self dolowercase dolowercase self removespace removespace self keepaccents keepaccents self vocabfile vocabfile self spmodel spm sentencepieceprocessorself spmodelkwargs self spmodel loadvocabfile used for whitespace normalization in input texts fmt off self whitespaces fmt on regular expression to remove nonprinting characters e g some unicode control chars in preprocessing self nonprintingcharactersre re compile f joinmapchr listrange0 9 listrange11 32 listrange127 160 160 173 8203 super init dolowercasedolowercase removespaceremovespace keepaccentskeepaccents bostokenbostoken eostokeneostoken unktokenunktoken padtokenpadtoken spmodelkwargsself spmodelkwargs kwargs copied from transformers models albert tokenizationalbert alberttokenizer getstate def getstateself state self dict copy statespmodel none return state copied from transformers models albert tokenizationalbert alberttokenizer setstate def setstateself d self dict d for backward compatibility if not hasattrself spmodelkwargs self spmodelkwargs self spmodel spm sentencepieceprocessorself spmodelkwargs self spmodel loadself vocabfile property copied from transformers models albert tokenizationalbert alberttokenizer vocabsize def vocabsizeself int return lenself spmodel def preprocesstextself text str str remove nonprinting characters text self nonprintingcharactersre sub text normalize whitespaces text joinchar if char not in self whitespaces else for char in text nfc unicode normalization text unicodedata normalizenfc text return text def tokenizeself text str kwargs liststr text self preprocesstexttext return self spmodel encodetext outtypestr def converttokentoidself token str int converts an index int to a token str using the vocab return self spmodel idtopieceindex staticmethod def cleanuptokenizationoutstring str str converts a sequence of tokens strings to a single string special tokens remain intact currentsubtokens outstring previsspecial false for token in tokens make sure that special tokens are not decoded using sentencepiece model if token in self allspecialtokens todo check if this is needed as it ensures that decodeencodedoc doc by adding extra whitespace in the decoded document if not previsspecial outstring outstring self spmodel decodecurrentsubtokens token previsspecial true currentsubtokens else currentsubtokens appendtoken previsspecial false outstring self spmodel decodecurrentsubtokens return outstring copied from transformers models albert tokenizationalbert alberttokenizer getvocab def getvocabself dictstr int vocab self convertidstotokensi i for i in rangeself vocabsize vocab updateself addedtokensencoder return vocab copied from transformers models albert tokenizationalbert alberttokenizer savevocabulary def savevocabularyself savedirectory str filenameprefix optionalstr none tuplestr if not os path isdirsavedirectory logger errorfvocabulary path savedirectory should be a directory return outvocabfile os path join savedirectory filenameprefix if filenameprefix else vocabfilesnamesvocabfile if os path abspathself vocabfile os path abspathoutvocabfile and os path isfileself vocabfile copyfileself vocabfile outvocabfile elif not os path isfileself vocabfile with openoutvocabfile wb as fi contentspiecemodel self spmodel serializedmodelproto fi writecontentspiecemodel return outvocabfile def encodefast self text unionstr liststr returntensors unionstr bool false unionlistint listlistint torch tensor if isinstancetext str text self preprocesstexttext tokenids self spmodel encodetext else text self preprocesstextt for t in text tokenids self spmodel encodetext if returntensors is true or returntensors pt tokenids torch tensortokenids return tokenids def decodefastself tokenids unionint listint str return self spmodel decodetokenids property def defaultchattemplateself logger warningonce nno chat template is defined for this tokenizer using the default template ffor the self class name class if the default is not appropriate for your model please set tokenizer chattemplate to an appropriate template see https huggingface codocstransformersmainchattemplating for more information n return eostoken bostoken for message in messages if message role user user message content else bot message content endif message text bostoken endfor bot construct an gptsw3 tokenizer based on sentencepiece https github com google sentencepiece this tokenizer inherits from pretrainedtokenizer which contains most of the main methods users should refer to this superclass for more information regarding those methods example usage python from transformers import gptsw3tokenizer tokenizer gptsw3tokenizer from_pretrained ai sweden gpt sw3 126m tokenizer svenska är kul input_ids 1814 377 3617 63504 args vocab_file str sentencepiece https github com google sentencepiece file generally has a spm extension that contains the vocabulary necessary to instantiate a tokenizer do_lower_case bool optional defaults to false whether or not to lowercase the input when tokenizing remove_space bool optional defaults to false whether or not to strip the text when tokenizing removing excess spaces before and after the string keep_accents bool optional defaults to false whether or not to keep accents when tokenizing pad_token str optional the token used for padding for example when batching sequences of different lengths if not provided will default to pad or unk depending on model size unk_token str optional the unknown token a token that is not in the vocabulary cannot be converted to an id and is set to be this token instead if not provided will default to unk eos_token str optional the end of sequence token seen during pretraining if not provided will default to endoftext bos_token str optional the beginning of sequence token that can be used for downstream task was not seen during pretraining if not provided will default to s or endoftext depending on model size sp_model_kwargs dict optional will be passed to the sentencepieceprocessor __init__ method the python wrapper for sentencepiece https github com google sentencepiece tree master python can be used among other things to set enable_sampling enable subword regularization nbest_size sampling parameters for unigram invalid for bpe dropout nbest_size 0 1 no sampling is performed nbest_size 1 samples from the nbest_size results nbest_size 0 assuming that nbest_size is infinite and samples from the all hypothesis lattice using forward filtering and backward sampling algorithm alpha smoothing parameter for unigram sampling and dropout probability of merge operations for bpe dropout attributes sp_model sentencepieceprocessor the sentencepiece processor that is used for every conversion string tokens and ids whitespaces set the whitespaces that are replaced in the whitespace normalization in preprocessing non_printing_characters_re pattern the compiled regular expression to remove non printing characters in preprocessing default definitions for our 2 tokenizer versions with none checks to enable proper testing used for whitespace normalization in input texts fmt off fmt on regular expression to remove non printing characters e g some unicode control chars in preprocessing copied from transformers models albert tokenization_albert alberttokenizer __getstate__ copied from transformers models albert tokenization_albert alberttokenizer __setstate__ for backward compatibility copied from transformers models albert tokenization_albert alberttokenizer vocab_size returns the preprocessed text this procedure is identical to what was used when training the tokenizer remove non printing characters normalize whitespaces nfc unicode normalization converts a token str to an id int using the vocab converts an index int to a token str using the vocab returns the input string this function is overridden to remove the default clean up converts a sequence of tokens strings to a single string special tokens remain intact make sure that special tokens are not decoded using sentencepiece model todo check if this is needed as it ensures that decode encode doc doc by adding extra whitespace in the decoded document copied from transformers models albert tokenization_albert alberttokenizer get_vocab copied from transformers models albert tokenization_albert alberttokenizer save_vocabulary encodes a text or batch of texts to token ids using preprocessing and the raw sp tokenizer this has reduced functionality but is often much faster does not handle special tokens correctly these can manually be added as ids afterwards does not support padding these can manually be added as ids afterwards use default huggingface tokenization methods for full functionality args text str or list str one or several text s to convert to token ids return_tensors str or bool returns pytorch tensors if set to true or pt returns list int list list int or torch tensor the encoded text s as token ids encodes a text or batch of texts to token ids using preprocessing and the raw sp tokenizer this has reduced functionality but is often much faster args token_ids int or list int encoded token or text as token id s returns str decoded text this chat template formats messages like an instant messenger chat log with user and bot strings preceding messages bos tokens are added between all messages
import os
import re
import unicodedata
from shutil import copyfile
from typing import 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


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

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

PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
    "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 GPTSw3Tokenizer(PreTrainedTokenizer):
    

    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    model_input_names = ["input_ids", "token_type_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        do_lower_case=False,
        remove_space=False,
        keep_accents=False,
        pad_token=None,
        unk_token=None,
        eos_token=None,
        bos_token=None,
        sp_model_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> None:
        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs

        name_or_path = 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"
            )
            name_or_path = "None"

        
        eos_token = "<|endoftext|>" if eos_token is None else eos_token
        unk_token = "<unk>" if unk_token is None else unk_token
        if "gpt-sw3-7b" in name_or_path:
            pad_token = unk_token if pad_token is None else pad_token
            bos_token = eos_token if bos_token is None else bos_token
        else:
            pad_token = "<pad>" if pad_token is None else pad_token
            bos_token = "<s>" if bos_token is None else bos_token

        self.do_lower_case = do_lower_case
        self.remove_space = remove_space
        self.keep_accents = keep_accents
        self.vocab_file = vocab_file

        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.Load(vocab_file)

        
        
        self.whitespaces = {" ", " ", " ", " ", " ", "　", " ", " ", " ", " ", "￼", ""}
        

        
        self.non_printing_characters_re = re.compile(
            f"[{''.join(map(chr, list(range(0, 9)) + list(range(11, 32)) + list(range(127, 160)) + [160, 173, 8203]))}]"
        )

        super().__init__(
            do_lower_case=do_lower_case,
            remove_space=remove_space,
            keep_accents=keep_accents,
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            pad_token=pad_token,
            sp_model_kwargs=self.sp_model_kwargs,
            **kwargs,
        )

    
    def __getstate__(self):
        state = self.__dict__.copy()
        state["sp_model"] = None
        return state

    
    def __setstate__(self, d):
        self.__dict__ = d

        
        if not hasattr(self, "sp_model_kwargs"):
            self.sp_model_kwargs = {}

        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.Load(self.vocab_file)

    @property
    
    def vocab_size(self) -> int:
        return len(self.sp_model)

    def preprocess_text(self, text: str) -> str:
        

        
        text = self.non_printing_characters_re.sub("", text)

        
        text = "".join([char if char not in self.whitespaces else " " for char in text])

        
        text = unicodedata.normalize("NFC", text)
        return text

    def _tokenize(self, text: str, **kwargs) -> List[str]:
        text = self.preprocess_text(text)
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token: str) -> int:
        
        return self.sp_model.PieceToId(token)

    def _convert_id_to_token(self, index: int) -> str:
        
        return self.sp_model.IdToPiece(index)

    @staticmethod
    def clean_up_tokenization(out_string: str) -> str:
        
        return out_string

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        
        current_sub_tokens = []
        out_string = ""
        prev_is_special = False
        for token in tokens:
            
            if token in self.all_special_tokens:
                
                if not prev_is_special:
                    out_string += " "

                out_string += self.sp_model.decode(current_sub_tokens) + token
                prev_is_special = True
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
                prev_is_special = False
        out_string += self.sp_model.decode(current_sub_tokens)

        return out_string

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

    
    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        if not os.path.isdir(save_directory):
            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
            return
        out_vocab_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        )

        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
            copyfile(self.vocab_file, out_vocab_file)
        elif not os.path.isfile(self.vocab_file):
            with open(out_vocab_file, "wb") as fi:
                content_spiece_model = self.sp_model.serialized_model_proto()
                fi.write(content_spiece_model)

        return (out_vocab_file,)

    def encode_fast(
        self, text: Union[str, List[str]], return_tensors: Union[str, bool] = False
    ) -> Union[List[int], List[List[int]], "torch.Tensor"]:
        

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

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

        return token_ids

    def decode_fast(self, token_ids: Union[int, List[int]]) -> str:
        

        return self.sp_model.decode(token_ids)

    @property
    def default_chat_template(self):
        
        logger.warning_once(
            "\nNo chat template is defined for this tokenizer - using the default template "
            f"for the {self.__class__.__name__} class. If the default is not appropriate for "
            "your model, please set `tokenizer.chat_template` to an appropriate template. "
            "See https://huggingface.co/docs/transformers/main/chat_templating for more information.\n"
        )
        return (
            "{{ eos_token }}{{ bos_token }}"
            "{% for message in messages %}"
            "{% if message['role'] == 'user' %}{{ 'User: ' + message['content']}}"
            "{% else %}{{ 'Bot: ' + message['content']}}{% endif %}"
            "{{ message['text'] }}{{ bos_token }}"
            "{% endfor %}"
            "Bot:"
        )