CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
tokenization_pegasus.py293 linesDownload Raw Back to pegasus
1# coding=utf-82# Copyright 2020 Google and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15import os16from shutil import copyfile17from typing import Any, Optional18 19import sentencepiece as spm20 21from ...tokenization_utils import AddedToken, PreTrainedTokenizer22from ...utils import logging23from ...utils.import_utils import requires24 25 26SPIECE_UNDERLINE = "▁"27 28VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}29 30 31logger = logging.get_logger(__name__)32 33 34# TODO ArthurZ refactor this to only use the added_tokens_encoder35 36 37@requires(backends=("sentencepiece",))38class PegasusTokenizer(PreTrainedTokenizer):39    r"""40    Construct a PEGASUS tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).41 42    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to43    this superclass for more information regarding those methods.44 45    Args:46        vocab_file (`str`):47            [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that48            contains the vocabulary necessary to instantiate a tokenizer.49        pad_token (`str`, *optional*, defaults to `"<pad>"`):50            The token used for padding, for example when batching sequences of different lengths.51        eos_token (`str`, *optional*, defaults to `"</s>"`):52            The end of sequence token.53 54            <Tip>55 56            When building a sequence using special tokens, this is not the token that is used for the end of sequence.57            The token used is the `sep_token`.58 59            </Tip>60 61        unk_token (`str`, *optional*, defaults to `"<unk>"`):62            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this63            token instead.64        mask_token (`str`, *optional*, defaults to `"<mask_2>"`):65            The token used for masking single token values. This is the token used when training this model with masked66            language modeling (MLM). This is the token that the PEGASUS encoder will try to predict during pretraining.67            It corresponds to *[MASK2]* in [PEGASUS: Pre-training with Extracted Gap-sentences for Abstractive68            Summarization](https://huggingface.co/papers/1912.08777).69        mask_token_sent (`str`, *optional*, defaults to `"<mask_1>"`):70            The token used for masking whole target sentences. This is the token used when training this model with gap71            sentences generation (GSG). This is the sentence that the PEGASUS decoder will try to predict during72            pretraining. It corresponds to *[MASK1]* in [PEGASUS: Pre-training with Extracted Gap-sentences for73            Abstractive Summarization](https://huggingface.co/papers/1912.08777).74        additional_special_tokens (`List[str]`, *optional*):75            Additional special tokens used by the tokenizer. If no additional_special_tokens are provided <mask_2> and76            <unk_2, ..., unk_102> are used as additional special tokens corresponding to the [original PEGASUS77            tokenizer](https://github.com/google-research/pegasus/blob/939830367bcf411193d2b5eca2f2f90f3f9260ca/pegasus/ops/pretrain_parsing_ops.cc#L66)78            that uses the tokens 2 - 104 only for pretraining79        sp_model_kwargs (`dict`, *optional*):80            Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for81            SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,82            to set:83 84            - `enable_sampling`: Enable subword regularization.85            - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.86 87              - `nbest_size = {0,1}`: No sampling is performed.88              - `nbest_size > 1`: samples from the nbest_size results.89              - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)90                using forward-filtering-and-backward-sampling algorithm.91 92            - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for93              BPE-dropout.94    """95 96    vocab_files_names = VOCAB_FILES_NAMES97    model_input_names = ["input_ids", "attention_mask"]98 99    def __init__(100        self,101        vocab_file,102        pad_token="<pad>",103        eos_token="</s>",104        unk_token="<unk>",105        mask_token="<mask_2>",106        mask_token_sent="<mask_1>",107        additional_special_tokens=None,108        offset=103,  # entries 2 - 104 are only used for pretraining109        sp_model_kwargs: Optional[dict[str, Any]] = None,110        **kwargs,111    ) -> None:112        self.offset = offset113        if additional_special_tokens is not None:114            if not isinstance(additional_special_tokens, list):115                raise TypeError(116                    f"additional_special_tokens should be of type {type(list)}, but is"117                    f" {type(additional_special_tokens)}"118                )119            additional_special_tokens_extended = (120                ([mask_token_sent] + additional_special_tokens)121                if mask_token_sent not in additional_special_tokens and mask_token_sent is not None122                else additional_special_tokens123            )124            # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken125            additional_special_tokens_extended += [126                f"<unk_{i}>" for i in range(len(additional_special_tokens_extended), self.offset - 1)127            ]128 129            if len(set(additional_special_tokens_extended)) != len(additional_special_tokens_extended):130                raise ValueError(131                    "Please make sure that the provided additional_special_tokens do not contain an incorrectly"132                    f" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."133                )134            additional_special_tokens = additional_special_tokens_extended135        else:136            additional_special_tokens_extended = []137            additional_special_tokens = [mask_token_sent] if mask_token_sent is not None else []138            additional_special_tokens += [f"<unk_{i}>" for i in range(2, self.offset)]139 140        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs141        self.mask_token_sent = mask_token_sent142        self.vocab_file = vocab_file143        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)144        self.sp_model.Load(vocab_file)145 146        _added_tokens_decoder = {147            0: AddedToken(str(pad_token), special=True),148            1: AddedToken(str(eos_token), special=True),149        }150 151        if self.mask_token_sent is not None:152            _added_tokens_decoder[2] = AddedToken(mask_token_sent, special=True)153            _added_tokens_decoder[3] = AddedToken(str(mask_token), special=True)154 155        for i in range(2, self.offset):156            _added_tokens_decoder[len(_added_tokens_decoder)] = AddedToken(f"<unk_{i}>", special=True)157 158        # Force update as we want to make sure vocab is enforced (same as fast)159        self._added_tokens_decoder = kwargs.pop("added_tokens_decoder", {})160        self._added_tokens_decoder.update(_added_tokens_decoder)161 162        super().__init__(163            eos_token=eos_token,164            unk_token=unk_token,165            mask_token=mask_token,166            pad_token=pad_token,167            mask_token_sent=mask_token_sent,168            offset=offset,169            additional_special_tokens=additional_special_tokens,170            sp_model_kwargs=self.sp_model_kwargs,171            **kwargs,172        )173 174    @property175    def vocab_size(self) -> int:176        return len(self.sp_model) + self.offset177 178    def get_vocab(self) -> dict[str, int]:179        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}180        vocab.update(self.added_tokens_encoder)181        return vocab182 183    def __getstate__(self):184        state = self.__dict__.copy()185        state["sp_model"] = None186        return state187 188    def __setstate__(self, d):189        self.__dict__ = d190 191        # for backward compatibility192        if not hasattr(self, "sp_model_kwargs"):193            self.sp_model_kwargs = {}194 195        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)196        self.sp_model.Load(self.vocab_file)197 198    def _tokenize(self, text: str) -> list[str]:199        """Take as input a string and return a list of strings (tokens) for words/sub-words"""200        return self.sp_model.encode(text, out_type=str)201 202    def _convert_token_to_id(self, token: str) -> int:203        """Converts a token (str) to an id using the vocab."""204        sp_id = self.sp_model.piece_to_id(token)205        return sp_id + self.offset206 207    def _convert_id_to_token(self, index: int) -> str:208        """Converts an index (integer) to a token (str) using the vocab."""209        if index < self.offset:210            return self.sp_model.IdToPiece(index)211        token = self.sp_model.IdToPiece(index - self.offset)212        return token213 214    def convert_tokens_to_string(self, tokens):215        """Converts a sequence of tokens (string) in a single string."""216        current_sub_tokens = []217        out_string = ""218        for token in tokens:219            # make sure that special tokens are not decoded using sentencepiece model220            if token in self.all_special_tokens:221                out_string += self.sp_model.decode(current_sub_tokens) + token222                current_sub_tokens = []223            else:224                current_sub_tokens.append(token)225        out_string += self.sp_model.decode(current_sub_tokens)226        return out_string.strip()227 228    def num_special_tokens_to_add(self, pair=False):229        """Just EOS"""230        return 1231 232    def _special_token_mask(self, seq):233        all_special_ids = set(self.all_special_ids)  # call it once instead of inside list comp234        all_special_ids.remove(self.unk_token_id)  # <unk> is only sometimes special235 236        return [1 if x in all_special_ids else 0 for x in seq]237 238    def get_special_tokens_mask(239        self, token_ids_0: list, token_ids_1: Optional[list] = None, already_has_special_tokens: bool = False240    ) -> list[int]:241        """Get list where entries are [1] if a token is [eos] or [pad] else 0."""242        if already_has_special_tokens:243            return self._special_token_mask(token_ids_0)244        elif token_ids_1 is None:245            return self._special_token_mask(token_ids_0) + [1]246        else:247            return self._special_token_mask(token_ids_0 + token_ids_1) + [1]248 249    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None) -> list[int]:250        """251        Build model inputs from a sequence or a pair of sequences for sequence classification tasks by concatenating252        and adding special tokens. A PEGASUS sequence has the following format, where `X` represents the sequence:253 254        - single sequence: `X </s>`255        - pair of sequences: `A B </s>` (not intended use)256 257        BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a258        separator.259 260        Args:261            token_ids_0 (`List[int]`):262                List of IDs to which the special tokens will be added.263            token_ids_1 (`List[int]`, *optional*):264                Optional second list of IDs for sequence pairs.265 266        Returns:267            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.268        """269        if token_ids_1 is None:270            return token_ids_0 + [self.eos_token_id]271        # We don't expect to process pairs, but leave the pair logic for API consistency272        return token_ids_0 + token_ids_1 + [self.eos_token_id]273 274    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:275        if not os.path.isdir(save_directory):276            logger.error(f"Vocabulary path ({save_directory}) should be a directory")277            return278        out_vocab_file = os.path.join(279            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]280        )281 282        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):283            copyfile(self.vocab_file, out_vocab_file)284        elif not os.path.isfile(self.vocab_file):285            with open(out_vocab_file, "wb") as fi:286                content_spiece_model = self.sp_model.serialized_model_proto()287                fi.write(content_spiece_model)288 289        return (out_vocab_file,)290 291 292__all__ = ["PegasusTokenizer"]293 
Aluode/PerceptionLabPortable · CoolFace