codingutf8 2021 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 tokenization classes for speech2text import json import os from pathlib import path from shutil import copyfile from typing import any dict list optional tuple union import sentencepiece from tokenizationutils import pretrainedtokenizer from utils import logging logger logging getloggername spieceunderline vocabfilesnames vocabfile vocab json spmfile sentencepiece bpe model pretrainedvocabfilesmap vocabfile facebooks2tsmalllibrispeechasr https huggingface cofacebooks2tsmalllibrispeechasrresolvemainvocab json spmfile facebooks2tsmalllibrispeechasr https huggingface cofacebooks2tsmalllibrispeechasrresolvemainsentencepiece bpe model maxmodelinputsizes facebooks2tsmalllibrispeechasr 1024 mustclangs pt fr ru nl ro it es de languages mustc mustclangs class speech2texttokenizerpretrainedtokenizer vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxmodelinputsizes maxmodelinputsizes modelinputnames inputids attentionmask prefixtokens listint def init self vocabfile spmfile bostokens eostokens padtokenpad unktokenunk douppercasefalse dolowercasefalse tgtlangnone langcodesnone additionalspecialtokensnone spmodelkwargs optionaldictstr any none kwargs none self spmodelkwargs if spmodelkwargs is none else spmodelkwargs self douppercase douppercase self dolowercase dolowercase self encoder loadjsonvocabfile self decoder v k for k v in self encoder items self spmfile spmfile self spmodel loadspmspmfile self spmodelkwargs if langcodes is not none self langcodes langcodes self langs languageslangcodes self langtokens flang lang for lang in self langs self langcodetoid lang self spmodel piecetoidflang lang for lang in self langs if additionalspecialtokens is not none additionalspecialtokens self langtokens additionalspecialtokens else additionalspecialtokens self langtokens self tgtlang tgtlang if tgtlang is not none else self langs0 self settgtlangspecialtokensself tgtlang else self langcodetoid super init bostokenbostoken eostokeneostoken unktokenunktoken padtokenpadtoken douppercasedouppercase dolowercasedolowercase tgtlangtgtlang langcodeslangcodes spmodelkwargsself spmodelkwargs additionalspecialtokensadditionalspecialtokens kwargs property def vocabsizeself int return lenself encoder def getvocabself dict vocab self encoder copy vocab updateself addedtokensencoder return vocab property def tgtlangself str return self tgtlang tgtlang setter def tgtlangself newtgtlang none self tgtlang newtgtlang self settgtlangspecialtokensnewtgtlang def settgtlangspecialtokensself tgtlang str none converts an index integer in a token str using the decoder return self decoder getindex self unktoken def converttokenstostringself tokens liststr str make sure that special tokens are not decoded using sentencepiece model build model inputs from a sequence by appending eostokenid if tokenids1 is none return self prefixtokens tokenids0 self eostokenid we don t expect to process pairs but leave the pair logic for api consistency return self prefixtokens tokenids0 tokenids1 self eostokenid def getspecialtokensmask self tokenids0 listint tokenids1 optionallistint none alreadyhasspecialtokens bool false listint if alreadyhasspecialtokens return super getspecialtokensmask tokenids0tokenids0 tokenids1tokenids1 alreadyhasspecialtokenstrue prefixones 1 lenself prefixtokens suffixones 1 if tokenids1 is none return prefixones 0 lentokenids0 suffixones return prefixones 0 lentokenids0 0 lentokenids1 suffixones def getstateself dict state self dict copy statespmodel none return state def setstateself d dict none self dict d for backward compatibility if not hasattrself spmodelkwargs self spmodelkwargs self spmodel loadspmself spmfile self spmodelkwargs def savevocabularyself savedirectory str filenameprefix optionalstr none tuplestr savedir pathsavedirectory assert savedir isdir fsavedirectory should be a directory vocabsavepath savedir filenameprefix if filenameprefix else self vocabfilesnamesvocabfile spmsavepath savedir filenameprefix if filenameprefix else self vocabfilesnamesspmfile savejsonself encoder vocabsavepath if os path abspathself spmfile os path abspathspmsavepath and os path isfileself spmfile copyfileself spmfile spmsavepath elif not os path isfileself spmfile with openspmsavepath wb as fi contentspiecemodel self spmodel serializedmodelproto fi writecontentspiecemodel return strvocabsavepath strspmsavepath def loadspmpath str spmodelkwargs dictstr any sentencepiece sentencepieceprocessor spm sentencepiece sentencepieceprocessorspmodelkwargs spm loadstrpath return spm def loadjsonpath str uniondict list with openpath r as f return json loadf def savejsondata path str none with openpath w as f json dumpdata f indent2 coding utf 8 2021 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 tokenization classes for speech2text construct an speech2text tokenizer this tokenizer inherits from pretrainedtokenizer which contains some of the main methods users should refer to the superclass for more information regarding such methods args vocab_file str file containing the vocabulary spm_file str path to the sentencepiece https github com google sentencepiece model file bos_token str optional defaults to s the beginning of sentence token eos_token str optional defaults to s the end of sentence token 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 pad_token str optional defaults to pad the token used for padding for example when batching sequences of different lengths do_upper_case bool optional defaults to false whether or not to uppercase the output when decoding do_lower_case bool optional defaults to false whether or not to lowercase the input when tokenizing tgt_lang str optional a string representing the target language 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 kwargs additional keyword arguments passed along to pretrainedtokenizer reset the special tokens to the target language setting prefix eos tgt_lang_code and suffix eos converts an index integer in a token str using the decoder converts a sequence of tokens strings for sub words in a single string make sure that special tokens are not decoded using sentencepiece model build model inputs from a sequence by appending eos_token_id we don t expect to process pairs but leave the pair logic for api consistency retrieve sequence ids from a token list that has no special tokens added this method is called when adding special tokens using the tokenizer prepare_for_model method args token_ids_0 list int list of ids token_ids_1 list int optional optional second list of ids for sequence pairs already_has_special_tokens bool optional defaults to false whether or not the token list is already formatted with special tokens for the model returns list int a list of integers in the range 0 1 1 for a special token 0 for a sequence token for backward compatibility
import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


logger = logging.get_logger(__name__)

SPIECE_UNDERLINE = "▁"

VOCAB_FILES_NAMES = {
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {
        "facebook/s2t-small-librispeech-asr": (
            "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/vocab.json"
        ),
    },
    "spm_file": {
        "facebook/s2t-small-librispeech-asr": (
            "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/sentencepiece.bpe.model"
        )
    },
}

MAX_MODEL_INPUT_SIZES = {
    "facebook/s2t-small-librispeech-asr": 1024,
}

MUSTC_LANGS = ["pt", "fr", "ru", "nl", "ro", "it", "es", "de"]

LANGUAGES = {"mustc": MUSTC_LANGS}


class Speech2TextTokenizer(PreTrainedTokenizer):
    

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

    prefix_tokens: List[int] = []

    def __init__(
        self,
        vocab_file,
        spm_file,
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
        unk_token="<unk>",
        do_upper_case=False,
        do_lower_case=False,
        tgt_lang=None,
        lang_codes=None,
        additional_special_tokens=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

        self.do_upper_case = do_upper_case
        self.do_lower_case = do_lower_case

        self.encoder = load_json(vocab_file)
        self.decoder = {v: k for k, v in self.encoder.items()}
        self.spm_file = spm_file
        self.sp_model = load_spm(spm_file, self.sp_model_kwargs)

        if lang_codes is not None:
            self.lang_codes = lang_codes
            self.langs = LANGUAGES[lang_codes]
            self.lang_tokens = [f"<lang:{lang}>" for lang in self.langs]
            self.lang_code_to_id = {lang: self.sp_model.PieceToId(f"<lang:{lang}>") for lang in self.langs}
            if additional_special_tokens is not None:
                additional_special_tokens = self.lang_tokens + additional_special_tokens
            else:
                additional_special_tokens = self.lang_tokens
            self._tgt_lang = tgt_lang if tgt_lang is not None else self.langs[0]

            self.set_tgt_lang_special_tokens(self._tgt_lang)
        else:
            self.lang_code_to_id = {}

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            pad_token=pad_token,
            do_upper_case=do_upper_case,
            do_lower_case=do_lower_case,
            tgt_lang=tgt_lang,
            lang_codes=lang_codes,
            sp_model_kwargs=self.sp_model_kwargs,
            additional_special_tokens=additional_special_tokens,
            **kwargs,
        )

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

    def get_vocab(self) -> Dict:
        vocab = self.encoder.copy()
        vocab.update(self.added_tokens_encoder)
        return vocab

    @property
    def tgt_lang(self) -> str:
        return self._tgt_lang

    @tgt_lang.setter
    def tgt_lang(self, new_tgt_lang) -> None:
        self._tgt_lang = new_tgt_lang
        self.set_tgt_lang_special_tokens(new_tgt_lang)

    def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None:
        
        lang_code_id = self.lang_code_to_id[tgt_lang]
        self.prefix_tokens = [lang_code_id]

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

    def _convert_token_to_id(self, token):
        return self.encoder.get(token, self.encoder[self.unk_token])

    def _convert_id_to_token(self, index: int) -> str:
        
        return self.decoder.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        
        current_sub_tokens = []
        out_string = ""
        for token in tokens:
            
            if token in self.all_special_tokens:
                decoded = self.sp_model.decode(current_sub_tokens)
                out_string += (decoded.upper() if self.do_upper_case else decoded) + token + " "
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
        decoded = self.sp_model.decode(current_sub_tokens)
        out_string += decoded.upper() if self.do_upper_case else decoded
        return out_string.strip()

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

    def get_special_tokens_mask(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
    ) -> List[int]:
        

        if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
            )

        prefix_ones = [1] * len(self.prefix_tokens)
        suffix_ones = [1]
        if token_ids_1 is None:
            return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
        return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones

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

    def __setstate__(self, d: Dict) -> None:
        self.__dict__ = d

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

        self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        save_dir = Path(save_directory)
        assert save_dir.is_dir(), f"{save_directory} should be a directory"
        vocab_save_path = save_dir / (
            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
        )
        spm_save_path = save_dir / (
            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
        )

        save_json(self.encoder, vocab_save_path)

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

        return (str(vocab_save_path), str(spm_save_path))


def load_spm(path: str, sp_model_kwargs: Dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
    spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
    spm.Load(str(path))
    return spm


def load_json(path: str) -> Union[Dict, List]:
    with open(path, "r") as f:
        return json.load(f)


def save_json(data, path: str) -> None:
    with open(path, "w") as f:
        json.dump(data, f, indent=2)