codingutf8 2020 google and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license todo arthurz refactor this to only use the addedtokensencoder vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxmodelinputsizes pretrainedpositionalembeddingssizes modelinputnames inputids attentionmask def init self vocabfile padtokenpad eostokens unktokenunk masktokenmask2 masktokensentmask1 additionalspecialtokensnone offset103 entries 2 104 are only used for pretraining spmodelkwargs optionaldictstr any none kwargs none self offset offset if additionalspecialtokens is not none if not isinstanceadditionalspecialtokens list raise typeerror fadditionalspecialtokens should be of type typelist but is f typeadditionalspecialtokens additionalspecialtokensextended masktokensent additionalspecialtokens if masktokensent not in additionalspecialtokens and masktokensent is not none else additionalspecialtokens fill additional tokens with unktoken102 in case not all additional tokens are already taken additionalspecialtokensextended funki for i in rangelenadditionalspecialtokensextended self offset 1 if lensetadditionalspecialtokensextended lenadditionalspecialtokensextended raise valueerror please make sure that the provided additionalspecialtokens do not contain an incorrectly f shifted list of unkx tokens found additionalspecialtokensextended additionalspecialtokens additionalspecialtokensextended else additionalspecialtokensextended additionalspecialtokens masktokensent if masktokensent is not none else additionalspecialtokens funki for i in range2 self offset self spmodelkwargs if spmodelkwargs is none else spmodelkwargs self masktokensent masktokensent self vocabfile vocabfile self spmodel spm sentencepieceprocessorself spmodelkwargs self spmodel loadvocabfile addedtokensdecoder 0 addedtokenstrpadtoken specialtrue 1 addedtokenstreostoken specialtrue if self masktokensent is not none addedtokensdecoder2 addedtokenmasktokensent specialtrue addedtokensdecoder3 addedtokenstrmasktoken specialtrue for i in range2 self offset addedtokensdecoderlenaddedtokensdecoder addedtokenfunki specialtrue force update as we want to make sure vocab is enforced same as fast self addedtokensdecoder kwargs popaddedtokensdecoder self addedtokensdecoder updateaddedtokensdecoder super init eostokeneostoken unktokenunktoken masktokenmasktoken padtokenpadtoken masktokensentmasktokensent offsetoffset additionalspecialtokensadditionalspecialtokens spmodelkwargsself spmodelkwargs kwargs property def vocabsizeself int return lenself spmodel self offset def getvocabself dictstr int vocab self convertidstotokensi i for i in rangeself vocabsize vocab updateself addedtokensencoder return vocab def getstateself state self dict copy statespmodel none return state def setstateself d self dict d for backward compatibility if not hasattrself spmodelkwargs self spmodelkwargs self spmodel spm sentencepieceprocessorself spmodelkwargs self spmodel loadself vocabfile def tokenizeself text str liststr converts a token str to an id using the vocab spid self spmodel piecetoidtoken return spid self offset def convertidtotokenself index int str converts a sequence of tokens string in a single string currentsubtokens outstring for token in tokens make sure that special tokens are not decoded using sentencepiece model if token in self allspecialtokens outstring self spmodel decodecurrentsubtokens token currentsubtokens else currentsubtokens appendtoken outstring self spmodel decodecurrentsubtokens return outstring strip def numspecialtokenstoaddself pairfalse get list where entries are 1 if a token is eos or pad else 0 if alreadyhasspecialtokens return self specialtokenmasktokenids0 elif tokenids1 is none return self specialtokenmasktokenids0 1 else return self specialtokenmasktokenids0 tokenids1 1 def buildinputswithspecialtokensself tokenids0 tokenids1none listint if tokenids1 is none return tokenids0 self eostokenid we don t expect to process pairs but leave the pair logic for api consistency return tokenids0 tokenids1 self eostokenid 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 coding utf 8 2020 google and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license todo arthurz refactor this to only use the added_tokens_encoder construct a pegasus 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 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 pad_token str optional defaults to pad the token used for padding for example when batching sequences of different lengths eos_token str optional defaults to s the end of sequence token tip when building a sequence using special tokens this is not the token that is used for the end of sequence the token used is the sep_token tip unk_token str optional defaults to unk 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 mask_token str optional defaults to mask_2 the token used for masking single token values this is the token used when training this model with masked language modeling mlm this is the token that the pegasus encoder will try to predict during pretraining it corresponds to mask2 in pegasus pre training with extracted gap sentences for abstractive summarization https arxiv org pdf 1912 08777 pdf mask_token_sent str optional defaults to mask_1 the token used for masking whole target sentences this is the token used when training this model with gap sentences generation gsg this is the sentence that the pegasus decoder will try to predict during pretraining it corresponds to mask1 in pegasus pre training with extracted gap sentences for abstractive summarization https arxiv org pdf 1912 08777 pdf additional_special_tokens list str optional additional special tokens used by the tokenizer if no additional_special_tokens are provided mask_2 and unk_2 unk_102 are used as additional special tokens corresponding to the original pegasus tokenizer https github com google research pegasus blob 939830367bcf411193d2b5eca2f2f90f3f9260ca pegasus ops pretrain_parsing_ops cc l66 that uses the tokens 2 104 only for pretraining 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 entries 2 104 are only used for pretraining fill additional tokens with unk_token_102 in case not all additional tokens are already taken force update as we want to make sure vocab is enforced same as fast for backward compatibility take as input a string and return a list of strings tokens for words sub words converts a token str to an id using the vocab converts an index integer to a token str using the vocab converts a sequence of tokens string in a single string make sure that special tokens are not decoded using sentencepiece model just eos call it once instead of inside list comp unk is only sometimes special get list where entries are 1 if a token is eos or pad else 0 build model inputs from a sequence or a pair of sequences for sequence classification tasks by concatenating and adding special tokens a pegasus sequence has the following format where x represents the sequence single sequence x s pair of sequences a b s not intended use bos is never used pairs of sequences are not the expected use case but they will be handled without a separator args token_ids_0 list int list of ids to which the special tokens will be added token_ids_1 list int optional optional second list of ids for sequence pairs returns list int list of input ids glossary input ids with the appropriate special tokens we don t expect to process pairs but leave the pair logic for api consistency
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


SPIECE_UNDERLINE = "▁"

VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"}
}

PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
    "google/pegasus-xsum": 512,
}


logger = logging.get_logger(__name__)



class PegasusTokenizer(PreTrainedTokenizer):
    r

    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", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        pad_token="<pad>",
        eos_token="</s>",
        unk_token="<unk>",
        mask_token="<mask_2>",
        mask_token_sent="<mask_1>",
        additional_special_tokens=None,
        offset=103,  
        sp_model_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> None:
        self.offset = offset
        if additional_special_tokens is not None:
            if not isinstance(additional_special_tokens, list):
                raise TypeError(
                    f"additional_special_tokens should be of type {type(list)}, but is"
                    f" {type(additional_special_tokens)}"
                )
            additional_special_tokens_extended = (
                ([mask_token_sent] + additional_special_tokens)
                if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                else additional_special_tokens
            )
            
            additional_special_tokens_extended += [
                f"<unk_{i}>" for i in range(len(additional_special_tokens_extended), self.offset - 1)
            ]

            if len(set(additional_special_tokens_extended)) != len(additional_special_tokens_extended):
                raise ValueError(
                    "Please make sure that the provided additional_special_tokens do not contain an incorrectly"
                    f" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."
                )
            additional_special_tokens = additional_special_tokens_extended
        else:
            additional_special_tokens_extended = []
            additional_special_tokens = [mask_token_sent] if mask_token_sent is not None else []
            additional_special_tokens += [f"<unk_{i}>" for i in range(2, self.offset)]

        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
        self.mask_token_sent = mask_token_sent
        self.vocab_file = vocab_file
        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.Load(vocab_file)

        _added_tokens_decoder = {
            0: AddedToken(str(pad_token), special=True),
            1: AddedToken(str(eos_token), special=True),
        }

        if self.mask_token_sent is not None:
            _added_tokens_decoder[2] = AddedToken(mask_token_sent, special=True)
            _added_tokens_decoder[3] = AddedToken(str(mask_token), special=True)

        for i in range(2, self.offset):
            _added_tokens_decoder[len(_added_tokens_decoder)] = AddedToken(f"<unk_{i}>", special=True)

        
        self._added_tokens_decoder = kwargs.pop("added_tokens_decoder", {})
        self._added_tokens_decoder.update(_added_tokens_decoder)

        super().__init__(
            eos_token=eos_token,
            unk_token=unk_token,
            mask_token=mask_token,
            pad_token=pad_token,
            mask_token_sent=mask_token_sent,
            offset=offset,
            additional_special_tokens=additional_special_tokens,
            sp_model_kwargs=self.sp_model_kwargs,
            **kwargs,
        )

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

    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 __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)

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

    def _convert_token_to_id(self, token: str) -> int:
        
        sp_id = self.sp_model.piece_to_id(token)
        return sp_id + self.offset

    def _convert_id_to_token(self, index: int) -> str:
        
        if index < self.offset:
            return self.sp_model.IdToPiece(index)
        token = self.sp_model.IdToPiece(index - self.offset)
        return token

    def convert_tokens_to_string(self, tokens):
        
        current_sub_tokens = []
        out_string = ""
        for token in tokens:
            
            if token in self.all_special_tokens:
                out_string += self.sp_model.decode(current_sub_tokens) + token
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
        out_string += self.sp_model.decode(current_sub_tokens)
        return out_string.strip()

    def num_special_tokens_to_add(self, pair=False):
        
        return 1

    def _special_token_mask(self, seq):
        all_special_ids = set(self.all_special_ids)  
        all_special_ids.remove(self.unk_token_id)  

        return [1 if x in all_special_ids else 0 for x in seq]

    def get_special_tokens_mask(
        self, token_ids_0: List, token_ids_1: Optional[List] = None, already_has_special_tokens: bool = False
    ) -> List[int]:
        
        if already_has_special_tokens:
            return self._special_token_mask(token_ids_0)
        elif token_ids_1 is None:
            return self._special_token_mask(token_ids_0) + [1]
        else:
            return self._special_token_mask(token_ids_0 + token_ids_1) + [1]

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None) -> List[int]:
        
        if token_ids_1 is None:
            return token_ids_0 + [self.eos_token_id]
        
        return token_ids_0 + token_ids_1 + [self.eos_token_id]

    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,)