codingutf8 2022 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 whisper import json import os import re from functools import lrucache from typing import list optional tuple import numpy as np from tokenizers import addedtoken pretokenizers processors from tokenizationutilsbase import batchencoding from tokenizationutilsfast import pretrainedtokenizerfast from utils import logging from englishnormalizer import basictextnormalizer englishtextnormalizer from tokenizationwhisper import languages taskids tolanguagecode whispertokenizer decodeasr logger logging getloggername vocabfilesnames vocabfile vocab json tokenizerfile tokenizer json mergesfile merges txt normalizerfile normalizer json pretrainedvocabfilesmap vocabfile openaiwhispertiny https huggingface coopenaiwhispertinyresolvemainvocab json openaiwhisperbase https huggingface coopenaiwhisperbaseresolvemainvocab json openaiwhispersmall https huggingface coopenaiwhispersmallresolvemainvocab json openaiwhispermedium https huggingface coopenaiwhispermediumresolvemainvocab json openaiwhisperlarge https huggingface coopenaiwhisperlargeresolvemainvocab json openaiwhispertiny en https huggingface coopenaiwhispertiny enresolvemainvocab json openaiwhisperbase en https huggingface coopenaiwhisperbase enresolvemainvocab json openaiwhispersmall en https huggingface coopenaiwhispersmall enresolvemainvocab json openaiwhispermedium en https huggingface coopenaiwhispermedium enresolvemainvocab json mergesfile openaiwhispertiny https huggingface coopenaiwhispertinyresolvemainmerges txt openaiwhisperbase https huggingface coopenaiwhisperbaseresolvemainmerges txt openaiwhispersmall https huggingface coopenaiwhispersmallresolvemainmerges txt openaiwhispermedium https huggingface coopenaiwhispermediumresolvemainmerges txt openaiwhisperlarge https huggingface coopenaiwhisperlargeresolvemainmerges txt openaiwhispertiny en https huggingface coopenaiwhispertiny enresolvemainmerges txt openaiwhisperbase en https huggingface coopenaiwhisperbase enresolvemainmerges txt openaiwhispersmall en https huggingface coopenaiwhispersmall enresolvemainmerges txt openaiwhispermedium en https huggingface coopenaiwhispermedium enresolvemainmerges txt tokenizerfile openaiwhispertiny https huggingface coopenaiwhispertinyresolvemaintokenizer json openaiwhisperbase https huggingface coopenaiwhisperbaseresolvemaintokenizer json openaiwhispersmall https huggingface coopenaiwhispersmallresolvemaintokenizer json openaiwhispermedium https huggingface coopenaiwhispermediumresolvemaintokenizer json openaiwhisperlarge https huggingface coopenaiwhisperlargeresolvemaintokenizer json openaiwhispertiny en https huggingface coopenaiwhispertiny enresolvemaintokenizer json openaiwhisperbase en https huggingface coopenaiwhisperbase enresolvemaintokenizer json openaiwhispersmall en https huggingface coopenaiwhispersmall enresolvemaintokenizer json openaiwhispermedium en https huggingface coopenaiwhispermedium enresolvemaintokenizer json pretrainedpositionalembeddingssizes openaiwhispertiny 1500 openaiwhisperbase 1500 openaiwhispersmall 1500 openaiwhispermedium 1500 openaiwhisperlarge 1500 openaiwhispertiny en 1500 openaiwhisperbase en 1500 openaiwhispersmall en 1500 openaiwhispermedium en 1500 class whispertokenizerfastpretrainedtokenizerfast vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxmodelinputsizes pretrainedpositionalembeddingssizes modelinputnames inputids attentionmask slowtokenizerclass whispertokenizer def init self vocabfilenone mergesfilenone normalizerfilenone tokenizerfilenone unktokenendoftext bostokenendoftext eostokenendoftext addprefixspacefalse languagenone tasknone predicttimestampsfalse kwargs bostoken addedtokenbostoken lstripfalse rstripfalse normalizedfalse specialtrue if isinstancebostoken str else bostoken eostoken addedtokeneostoken lstripfalse rstripfalse normalizedfalse specialtrue if isinstanceeostoken str else eostoken unktoken addedtokenunktoken lstripfalse rstripfalse normalizedfalse specialtrue if isinstanceunktoken str else unktoken super init vocabfile mergesfile tokenizerfiletokenizerfile unktokenunktoken bostokenbostoken eostokeneostoken addprefixspaceaddprefixspace kwargs self addbostoken kwargs popaddbostoken false pretokstate json loadsself backendtokenizer pretokenizer getstate if pretokstate getaddprefixspace addprefixspace addprefixspace pretokclass getattrpretokenizers pretokstate poptype pretokstateaddprefixspace addprefixspace self backendtokenizer pretokenizer pretokclasspretokstate if normalizerfile is not none with opennormalizerfile encodingutf8 as vocabhandle self englishspellingnormalizer json loadvocabhandle else self englishspellingnormalizer none self addprefixspace addprefixspace self timestamppat re compilerd d self language language self task task self predicttimestamps predicttimestamps copied from transformers models gpt2 tokenizationgpt2fast gpt2tokenizerfast batchencodeplus def batchencodeplusself args kwargs batchencoding issplitintowords kwargs getissplitintowords false assert self addprefixspace or not issplitintowords fyou need to instantiate self class name with addprefixspacetrue to use it with pretokenized inputs return super batchencodeplusargs kwargs copied from transformers models gpt2 tokenizationgpt2fast gpt2tokenizerfast encodeplus def encodeplusself args kwargs batchencoding issplitintowords kwargs getissplitintowords false assert self addprefixspace or not issplitintowords fyou need to instantiate self class name with addprefixspacetrue to use it with pretokenized inputs return super encodeplusargs kwargs copied from transformers models whisper tokenizationwhisper whispertokenizer decodewithtimestamps def decodewithtimestampsself tokenids skipspecialtokensfalse timeprecision0 02 str timestampbegin self allspecialids1 1 outputs for token in tokenids if token timestampbegin timestamp ftoken timestampbegin timeprecision 2f outputs appendtimestamp outputs append else outputs1 appendtoken outputs s if isinstances str else self decodes skipspecialtokensskipspecialtokens for s in outputs return joinoutputs copied from transformers models whisper tokenizationwhisper whispertokenizer computeoffsets def computeoffsetsself tokenids timeprecision0 02 offsets tokenids np arraytokenids if tokenids shape0 1 and lentokenids shape 1 raise valueerrorcan only process a single input at a time timestampbegin self allspecialids1 1 timestamptokens tokenids timestampbegin consecutive np wheretimestamptokens 1 timestamptokens1 0 1 if consecutive shape0 0 and timestamptokens sum 1 either there are no timestamps or there are no consecutive ones return elif np wheretimestamptokens01 1 not in consecutive we add the final timestamp if it is not already in the list consecutive np appendconsecutive np wheretimestamptokens01 1 lastslice np wheretimestamptokens00 for currentslice in consecutive slicedtokens tokenidslastslice currentslice if lenslicedtokens 1 starttimestampposition slicedtokens0 item timestampbegin endtimestampposition slicedtokens1 item timestampbegin strip timestamp tokens from the text output slicedtokens self preprocesstokenidsslicedtokens text self decodeslicedtokens text self filtertimestampidstext offsets append text text timestamp starttimestampposition timeprecision endtimestampposition timeprecision lastslice currentslice return offsets lrucache copied from transformers models whisper tokenizationwhisper whispertokenizer timestampids def timestampidsself timeprecision0 02 return self converttokenstoids 2f i timeprecision for i in range1500 1 copied from transformers models whisper tokenizationwhisper whispertokenizer preprocesstokenids def preprocesstokenidsself tokenids skipspecialtokens bool false if skipspecialtokens prompttokenid self converttokenstoidsstartofprev decoderstarttokenid self converttokenstoidsstartoftranscript tokenids self stripprompttokenids prompttokenid decoderstarttokenid return tokenids copied from transformers models whisper tokenizationwhisper whispertokenizer filtertimestampids def filtertimestampidsself tokenids return re subself timestamppat tokenids copied from transformers models whisper tokenizationwhisper whispertokenizer decode def decode self tokenids skipspecialtokens bool false cleanuptokenizationspaces bool none outputoffsets bool false timeprecision0 02 decodewithtimestamps bool false normalize bool false basicnormalize bool false removediacritics bool false kwargs str filteredids self preprocesstokenids tokenids skipspecialtokensskipspecialtokens text super decode filteredids skipspecialtokensskipspecialtokens cleanuptokenizationspacescleanuptokenizationspaces normalizenormalize basicnormalizebasicnormalize removediacriticsremovediacritics kwargs if decodewithtimestamps legacy method to decode timestamps when not included in the tokenizer vocabulary text self decodewithtimestamps filteredids timeprecisiontimeprecision skipspecialtokensskipspecialtokens else text self filtertimestampidstext retrieve offsets if outputoffsets offsets self computeoffsetstokenids timeprecisiontimeprecision return text text offsets offsets return text def decode self args normalize bool false basicnormalize bool false removediacritics bool false kwargs str text super decodeargs kwargs if normalize cleantext self normalizetext return cleantext elif basicnormalize cleantext self basicnormalizetext removediacriticsremovediacritics return cleantext else return text copied from transformers models whisper tokenizationwhisper whispertokenizer normalize def normalizeself text normalizer englishtextnormalizerself englishspellingnormalizer return normalizertext staticmethod copied from transformers models whisper tokenizationwhisper whispertokenizer basicnormalize def basicnormalizetext removediacriticsfalse normalizer basictextnormalizerremovediacriticsremovediacritics return normalizertext def savevocabularyself savedirectory str filenameprefix optionalstr none tuplestr files self tokenizer model savesavedirectory namefilenameprefix normalizerfile os path join savedirectory filenameprefix if filenameprefix else vocabfilesnamesnormalizerfile if self englishspellingnormalizer is not none with opennormalizerfile w encodingutf8 as f f write json dumpsself englishspellingnormalizer indent2 sortkeystrue ensureasciifalse n return tuplefiles normalizerfile def setprefixtokensself language str none task str none predicttimestamps bool none self language language if language is not none else self language self task task if task is not none else self task self predicttimestamps predicttimestamps if predicttimestamps is not none else self predicttimestamps prefixtokenids self prefixtokens prefixes self convertidstotokensprefixtokenids eos self eostoken eostokenid self eostokenid prefixtemplate joinftoken 0 for token in prefixes self backendtokenizer postprocessor processors templateprocessing singlefprefixtemplate a 0 eos 0 pairfprefixtemplate a 0 b 1 eos 1 specialtokens eos eostokenid zipprefixes prefixtokenids property copied from transformers models whisper tokenizationwhisper whispertokenizer prefixtokens def prefixtokensself listint bostokenid self converttokenstoidsstartoftranscript translatetokenid self converttokenstoidstranslate transcribetokenid self converttokenstoidstranscribe notimestampstokenid self converttokenstoidsnotimestamps langs tuplelanguages keys if self language is not none self language self language lower if self language in tolanguagecode languageid tolanguagecodeself language elif self language in tolanguagecode values languageid self language else islanguagecode lenself language 2 raise valueerror funsupported language self language language should be one of f listtolanguagecode values if islanguagecode else listtolanguagecode keys if self task is not none if self task not in taskids raise valueerrorfunsupported task self task task should be in taskids bossequence bostokenid if self language is not none bossequence appendbostokenid 1 langs indexlanguageid if self task is not none bossequence appendtranscribetokenid if self task transcribe else translatetokenid if not self predicttimestamps bossequence appendnotimestampstokenid return bossequence copied from transformers models whisper tokenizationwhisper whispertokenizer buildinputswithspecialtokens def buildinputswithspecialtokensself tokenids0 tokenids1none listint we don t expect to process pairs but leave the pair logic for api consistency copied from transformers models whisper tokenizationwhisper whispertokenizer getspecialtokensmask retrieve sequence ids from a token list that has no special tokens added this method is called when adding special tokens using the tokenizer prepareformodel method args tokenids0 listint list of ids tokenids1 listint optional optional second list of ids for sequence pairs alreadyhasspecialtokens bool optional defaults to false whether or not the token list is already formatted with special tokens for the model returns listint a list of integers in the range 0 1 1 for a special token 0 for a sequence token copied from transformers models gpt2 tokenizationgpt2 gpt2tokenizer defaultchattemplate a simple chat template that ignores role information and just concatenates messages with eos tokens copied from transformers models whisper tokenizationwhisper whispertokenizer getdecoderpromptids prefix tokens are of the form startoftranscript langid task notimestamps we don t want to force the bos token at position 1 as this is the starting token when we generate so we slice the prefix tokens to langid task notimestamps to get the forced tokens copied from transformers models whisper tokenizationwhisper whispertokenizer getpromptids converts prompt text to ids that can be passed to whisperforconditionalgeneration generate batchencoding selfstartofprev text strip addspecialtokensfalse check for special tokens prompttextids batchencodinginputids1 specialtokenid nextx for x in prompttextids if x self allspecialids0 none if specialtokenid is not none token self convertidstotokensspecialtokenid raise valueerrorfencountered text in the prompt corresponding to disallowed special token token batchencoding converttotensorstensortypereturntensors return batchencodinginputids staticmethod copied from transformers models whisper tokenizationwhisper whispertokenizer stripprompt def stripprompttokenids listint prompttokenid int decoderstarttokenid int hasprompt isinstancetokenids list and tokenids and tokenids0 prompttokenid if hasprompt if decoderstarttokenid in tokenids return tokenidstokenids indexdecoderstarttokenid else return return tokenids coding utf 8 2022 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 whisper construct a fast whisper tokenizer backed by huggingface s tokenizers library this tokenizer inherits from pretrainedtokenizerfast which contains most of the main methods users should refer to this superclass for more information regarding those methods args vocab_file str optional path to the vocabulary file merges_file str optional path to the merges file normalizer_file str optional path to the normalizer_file file tokenizer_file str optional path to tokenizers https github com huggingface tokenizers file generally has a json extension that contains everything needed to load the tokenizer unk_token str optional defaults to endoftext 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 bos_token str optional defaults to endoftext the beginning of sequence token the decoder_start_token_id is used to set the first token as startoftranscript when generating eos_token str optional defaults to endoftext the end of sequence token add_prefix_space bool optional defaults to false whether or not to add an initial space to the input this allows to treat the leading word just as any other word whisper tokenizer detect beginning of words by the preceding space language str optional the language of the transcription text the corresponding language id token is appended to the start of the sequence for multilingual speech recognition and speech translation tasks e g for spanish the token es is appended to the start of sequence this should be used for multilingual fine tuning only task str optional task identifier to append at the start of sequence if any this should be used for mulitlingual fine tuning with transcribe for speech recognition and translate for speech translation predict_timestamps bool optional defaults to false whether to omit the notimestamps token at the start of the sequence copied from transformers models gpt2 tokenization_gpt2_fast gpt2tokenizerfast _batch_encode_plus copied from transformers models gpt2 tokenization_gpt2_fast gpt2tokenizerfast _encode_plus copied from transformers models whisper tokenization_whisper whispertokenizer _decode_with_timestamps timestamp tokens are above the special tokens id range and are ignored by decode this method decodes given tokens with timestamps tokens annotated e g 1 08 copied from transformers models whisper tokenization_whisper whispertokenizer _compute_offsets compute offsets for a given tokenized input args token_ids union int list int np ndarray torch tensor tf tensor list of tokenized input ids can be obtained using the __call__ method time_precision float optional defaults to 0 02 the time ratio to convert from token to time either there are no timestamps or there are no consecutive ones we add the final timestamp if it is not already in the list strip timestamp tokens from the text output copied from transformers models whisper tokenization_whisper whispertokenizer timestamp_ids compute the timestamp token ids for a given precision and save to least recently used lru cache args time_precision float optional defaults to 0 02 the time ratio to convert from token to time copied from transformers models whisper tokenization_whisper whispertokenizer _preprocess_token_ids pre process the token ids for decoding by removing the prompt tokens ids and timestamp token ids args token_ids union int list int np ndarray torch tensor tf tensor list of tokenized input ids typically obtained using the __call__ method of the tokenizer skip_special_tokens bool optional defaults to false whether or not to remove special tokens from the token ids if true the prompt token ids will be removed copied from transformers models whisper tokenization_whisper whispertokenizer _filter_timestamp_ids copied from transformers models whisper tokenization_whisper whispertokenizer decode converts a sequence of ids in a string using the tokenizer and vocabulary with options to remove special tokens and clean up tokenization spaces similar to doing self convert_tokens_to_string self convert_ids_to_tokens token_ids args token_ids union int list int np ndarray torch tensor tf tensor list of tokenized input ids can be obtained using the __call__ method skip_special_tokens bool optional defaults to false whether or not to remove special tokens in the decoding clean_up_tokenization_spaces bool optional whether or not to clean up the tokenization spaces if none will default to self clean_up_tokenization_spaces available in the tokenizer_config output_offsets bool optional defaults to false whether or not to output the offsets of the tokens this should only be set if the model predicted timestamps time_precision float optional defaults to 0 02 the time ratio to convert from token to time decode_with_timestamps bool optional defaults to false whether or not to decode with timestamps included in the raw text normalize bool optional defaults to false whether or not to apply the english text normalizer to the decoded text only applicable when the target text is in english otherwise the basic text normalizer should be applied basic_normalize bool optional defaults to false whether or not to apply the basic text normalizer to the decoded text applicable to multilingual target text remove_diacritics bool optional defaults to false whether or not to remove diacritics when applying the basic text normalizer removing diacritics may destroy information in the decoded text hence it should be used with caution kwargs additional keyword arguments optional will be passed to the underlying model specific decode method returns str the decoded sentence legacy method to decode timestamps when not included in the tokenizer vocabulary retrieve offsets copied from transformers models whisper tokenization_whisper whispertokenizer _normalize normalize a given string using the englishtextnormalizer class which preforms commons transformation on english text copied from transformers models whisper tokenization_whisper whispertokenizer _basic_normalize normalize a given string using the basictextnormalizer class which preforms commons transformation on multilingual text override the prefix tokens appended to the start of the label sequence this method can be used standalone to update the prefix tokens as required when fine tuning example python instantiate the tokenizer and set the prefix token to spanish tokenizer whispertokenizerfast from_pretrained openai whisper tiny language spanish now switch the prefix token from spanish to french tokenizer set_prefix_tokens language french args language str optional defaults to none the language of the transcription text task str optional defaults to none task identifier to append at the start of sequence if any predict_timestamps bool optional defaults to none whether to omit the notimestamps token at the start of the sequence copied from transformers models whisper tokenization_whisper whispertokenizer prefix_tokens copied from transformers models whisper tokenization_whisper whispertokenizer build_inputs_with_special_tokens 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 copied from transformers models whisper tokenization_whisper whispertokenizer get_special_tokens_mask 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 copied from transformers models gpt2 tokenization_gpt2 gpt2tokenizer default_chat_template a simple chat template that ignores role information and just concatenates messages with eos tokens copied from transformers models whisper tokenization_whisper whispertokenizer get_decoder_prompt_ids prefix tokens are of the form startoftranscript lang_id task notimestamps we don t want to force the bos token at position 1 as this is the starting token when we generate so we slice the prefix tokens to lang_id task notimestamps to get the forced tokens copied from transformers models whisper tokenization_whisper whispertokenizer get_prompt_ids converts prompt text to ids that can be passed to whisperforconditionalgeneration generate check for special tokens copied from transformers models whisper tokenization_whisper whispertokenizer _strip_prompt
import json
import os
import re
from functools import lru_cache
from typing import List, Optional, Tuple

import numpy as np
from tokenizers import AddedToken, pre_tokenizers, processors

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .english_normalizer import BasicTextNormalizer, EnglishTextNormalizer
from .tokenization_whisper import LANGUAGES, TASK_IDS, TO_LANGUAGE_CODE, WhisperTokenizer, _decode_asr


logger = logging.get_logger(__name__)

VOCAB_FILES_NAMES = {
    "vocab_file": "vocab.json",
    "tokenizer_file": "tokenizer.json",
    "merges_file": "merges.txt",
    "normalizer_file": "normalizer.json",
}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {
        "openai/whisper-tiny": "https://huggingface.co/openai/whisper-tiny/resolve/main/vocab.json",
        "openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/vocab.json",
        "openai/whisper-small": "https://huggingface.co/openai/whisper-small/resolve/main/vocab.json",
        "openai/whisper-medium": "https://huggingface.co/openai/whisper-medium/resolve/main/vocab.json",
        "openai/whisper-large": "https://huggingface.co/openai/whisper-large/resolve/main/vocab.json",
        "openai/whisper-tiny.en": "https://huggingface.co/openai/whisper-tiny.en/resolve/main/vocab.json",
        "openai/whisper-base.en": "https://huggingface.co/openai/whisper-base.en/resolve/main/vocab.json",
        "openai/whisper-small.en": "https://huggingface.co/openai/whisper-small.en/resolve/main/vocab.json",
        "openai/whisper-medium.en": "https://huggingface.co/openai/whisper-medium.en/resolve/main/vocab.json",
    },
    "merges_file": {
        "openai/whisper-tiny": "https://huggingface.co/openai/whisper-tiny/resolve/main/merges.txt",
        "openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/merges.txt",
        "openai/whisper-small": "https://huggingface.co/openai/whisper-small/resolve/main/merges.txt",
        "openai/whisper-medium": "https://huggingface.co/openai/whisper-medium/resolve/main/merges.txt",
        "openai/whisper-large": "https://huggingface.co/openai/whisper-large/resolve/main/merges.txt",
        "openai/whisper-tiny.en": "https://huggingface.co/openai/whisper-tiny.en/resolve/main/merges.txt",
        "openai/whisper-base.en": "https://huggingface.co/openai/whisper-base.en/resolve/main/merges.txt",
        "openai/whisper-small.en": "https://huggingface.co/openai/whisper-small.en/resolve/main/merges.txt",
        "openai/whisper-medium.en": "https://huggingface.co/openai/whisper-medium.en/resolve/main/merges.txt",
    },
    "tokenizer_file": {
        "openai/whisper-tiny": "https://huggingface.co/openai/whisper-tiny/resolve/main/tokenizer.json",
        "openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/tokenizer.json",
        "openai/whisper-small": "https://huggingface.co/openai/whisper-small/resolve/main/tokenizer.json",
        "openai/whisper-medium": "https://huggingface.co/openai/whisper-medium/resolve/main/tokenizer.json",
        "openai/whisper-large": "https://huggingface.co/openai/whisper-large/resolve/main/tokenizer.json",
        "openai/whisper-tiny.en": "https://huggingface.co/openai/whisper-tiny.en/resolve/main/tokenizer.json",
        "openai/whisper-base.en": "https://huggingface.co/openai/whisper-base.en/resolve/main/tokenizer.json",
        "openai/whisper-small.en": "https://huggingface.co/openai/whisper-small.en/resolve/main/tokenizer.json",
        "openai/whisper-medium.en": "https://huggingface.co/openai/whisper-medium.en/resolve/main/tokenizer.json",
    },
}

PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
    "openai/whisper-tiny": 1500,
    "openai/whisper-base": 1500,
    "openai/whisper-small": 1500,
    "openai/whisper-medium": 1500,
    "openai/whisper-large": 1500,
    "openai/whisper-tiny.en": 1500,
    "openai/whisper-base.en": 1500,
    "openai/whisper-small.en": 1500,
    "openai/whisper-medium.en": 1500,
}


class WhisperTokenizerFast(PreTrainedTokenizerFast):
    

    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"]
    slow_tokenizer_class = WhisperTokenizer

    def __init__(
        self,
        vocab_file=None,
        merges_file=None,
        normalizer_file=None,
        tokenizer_file=None,
        unk_token="<|endoftext|>",
        bos_token="<|endoftext|>",
        eos_token="<|endoftext|>",
        add_prefix_space=False,
        language=None,
        task=None,
        predict_timestamps=False,
        **kwargs,
    ):
        bos_token = (
            AddedToken(bos_token, lstrip=False, rstrip=False, normalized=False, special=True)
            if isinstance(bos_token, str)
            else bos_token
        )
        eos_token = (
            AddedToken(eos_token, lstrip=False, rstrip=False, normalized=False, special=True)
            if isinstance(eos_token, str)
            else eos_token
        )
        unk_token = (
            AddedToken(unk_token, lstrip=False, rstrip=False, normalized=False, special=True)
            if isinstance(unk_token, str)
            else unk_token
        )

        super().__init__(
            vocab_file,
            merges_file,
            tokenizer_file=tokenizer_file,
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            add_prefix_space=add_prefix_space,
            **kwargs,
        )

        self.add_bos_token = kwargs.pop("add_bos_token", False)

        pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__())
        if pre_tok_state.get("add_prefix_space", add_prefix_space) != add_prefix_space:
            pre_tok_class = getattr(pre_tokenizers, pre_tok_state.pop("type"))
            pre_tok_state["add_prefix_space"] = add_prefix_space
            self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state)

        if normalizer_file is not None:
            with open(normalizer_file, encoding="utf-8") as vocab_handle:
                self.english_spelling_normalizer = json.load(vocab_handle)
        else:
            self.english_spelling_normalizer = None

        self.add_prefix_space = add_prefix_space
        self.timestamp_pat = re.compile(r"<\|(\d+\.\d+)\|>")

        self.language = language
        self.task = task
        self.predict_timestamps = predict_timestamps

    
    def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
        is_split_into_words = kwargs.get("is_split_into_words", False)
        assert self.add_prefix_space or not is_split_into_words, (
            f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
            "to use it with pretokenized inputs."
        )

        return super()._batch_encode_plus(*args, **kwargs)

    
    def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
        is_split_into_words = kwargs.get("is_split_into_words", False)

        assert self.add_prefix_space or not is_split_into_words, (
            f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
            "to use it with pretokenized inputs."
        )

        return super()._encode_plus(*args, **kwargs)

    
    def _decode_with_timestamps(self, token_ids, skip_special_tokens=False, time_precision=0.02) -> str:
        
        timestamp_begin = self.all_special_ids[-1] + 1
        outputs = [[]]
        for token in token_ids:
            if token >= timestamp_begin:
                timestamp = f"<|{(token - timestamp_begin) * time_precision:.2f}|>"
                outputs.append(timestamp)
                outputs.append([])
            else:
                outputs[-1].append(token)
        outputs = [
            s if isinstance(s, str) else self.decode(s, skip_special_tokens=skip_special_tokens) for s in outputs
        ]
        return "".join(outputs)

    
    def _compute_offsets(self, token_ids, time_precision=0.02):
        
        offsets = []
        token_ids = np.array(token_ids)
        if token_ids.shape[0] > 1 and len(token_ids.shape) > 1:
            raise ValueError("Can only process a single input at a time")
        timestamp_begin = self.all_special_ids[-1] + 1
        timestamp_tokens = token_ids >= timestamp_begin

        consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1
        if consecutive.shape[0] == 0 and timestamp_tokens.sum() <= 1:
            
            return []
        elif np.where(timestamp_tokens)[0][-1] + 1 not in consecutive:
            
            consecutive = np.append(consecutive, np.where(timestamp_tokens)[0][-1] + 1)

        last_slice = np.where(timestamp_tokens)[0][0]
        for current_slice in consecutive:
            sliced_tokens = token_ids[last_slice:current_slice]
            if len(sliced_tokens) > 1:
                start_timestamp_position = sliced_tokens[0].item() - timestamp_begin
                end_timestamp_position = sliced_tokens[-1].item() - timestamp_begin
                
                sliced_tokens = self._preprocess_token_ids(sliced_tokens)
                text = self._decode(sliced_tokens)
                text = self._filter_timestamp_ids(text)
                offsets.append(
                    {
                        "text": text,
                        "timestamp": (
                            start_timestamp_position * time_precision,
                            end_timestamp_position * time_precision,
                        ),
                    }
                )
            last_slice = current_slice

        return offsets

    @lru_cache
    
    def timestamp_ids(self, time_precision=0.02):
        
        return self.convert_tokens_to_ids([("<|%.2f|>" % (i * time_precision)) for i in range(1500 + 1)])

    
    def _preprocess_token_ids(self, token_ids, skip_special_tokens: bool = False):
        
        if skip_special_tokens:
            prompt_token_id = self.convert_tokens_to_ids("<|startofprev|>")
            decoder_start_token_id = self.convert_tokens_to_ids("<|startoftranscript|>")
            token_ids = self._strip_prompt(token_ids, prompt_token_id, decoder_start_token_id)

        return token_ids

    
    def _filter_timestamp_ids(self, token_ids):
        return re.sub(self.timestamp_pat, "", token_ids)

    
    def decode(
        self,
        token_ids,
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        output_offsets: bool = False,
        time_precision=0.02,
        decode_with_timestamps: bool = False,
        normalize: bool = False,
        basic_normalize: bool = False,
        remove_diacritics: bool = False,
        **kwargs,
    ) -> str:
        
        filtered_ids = self._preprocess_token_ids(
            token_ids,
            skip_special_tokens=skip_special_tokens,
        )

        text = super().decode(
            filtered_ids,
            skip_special_tokens=skip_special_tokens,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            normalize=normalize,
            basic_normalize=basic_normalize,
            remove_diacritics=remove_diacritics,
            **kwargs,
        )
        if decode_with_timestamps:
            
            text = self._decode_with_timestamps(
                filtered_ids, time_precision=time_precision, skip_special_tokens=skip_special_tokens
            )
        else:
            text = self._filter_timestamp_ids(text)

        
        if output_offsets:
            offsets = self._compute_offsets(token_ids, time_precision=time_precision)
            return {"text": text, "offsets": offsets}
        return text

    def _decode(
        self, *args, normalize: bool = False, basic_normalize: bool = False, remove_diacritics: bool = False, **kwargs
    ) -> str:
        text = super()._decode(*args, **kwargs)

        if normalize:
            clean_text = self._normalize(text)
            return clean_text
        elif basic_normalize:
            clean_text = self._basic_normalize(text, remove_diacritics=remove_diacritics)
            return clean_text
        else:
            return text

    
    def _normalize(self, text):
        
        normalizer = EnglishTextNormalizer(self.english_spelling_normalizer)
        return normalizer(text)

    @staticmethod
    
    def _basic_normalize(text, remove_diacritics=False):
        
        normalizer = BasicTextNormalizer(remove_diacritics=remove_diacritics)
        return normalizer(text)

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
        files = self._tokenizer.model.save(save_directory, name=filename_prefix)

        normalizer_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["normalizer_file"]
        )

        if self.english_spelling_normalizer is not None:
            with open(normalizer_file, "w", encoding="utf-8") as f:
                f.write(
                    json.dumps(self.english_spelling_normalizer, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
                )

        return tuple(files) + (normalizer_file,)

    def set_prefix_tokens(self, language: str = None, task: str = None, predict_timestamps: bool = None):
        
        self.language = language if language is not None else self.language
        self.task = task if task is not None else self.task
        self.predict_timestamps = predict_timestamps if predict_timestamps is not None else self.predict_timestamps

        prefix_token_ids = self.prefix_tokens
        prefixes = self.convert_ids_to_tokens(prefix_token_ids)
        eos = self.eos_token
        eos_token_id = self.eos_token_id
        prefix_template = " ".join([f"{token}:0" for token in prefixes])
        self.backend_tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{prefix_template} $A:0 {eos}:0",
            pair=f"{prefix_template} $A:0 $B:1 {eos}:1",
            special_tokens=[
                (eos, eos_token_id),
                *zip(prefixes, prefix_token_ids),
            ],
        )

    @property
    
    def prefix_tokens(self) -> List[int]:
        bos_token_id = self.convert_tokens_to_ids("<|startoftranscript|>")
        translate_token_id = self.convert_tokens_to_ids("<|translate|>")
        transcribe_token_id = self.convert_tokens_to_ids("<|transcribe|>")
        notimestamps_token_id = self.convert_tokens_to_ids("<|notimestamps|>")
        langs = tuple(LANGUAGES.keys())

        if self.language is not None:
            self.language = self.language.lower()
            if self.language in TO_LANGUAGE_CODE:
                language_id = TO_LANGUAGE_CODE[self.language]
            elif self.language in TO_LANGUAGE_CODE.values():
                language_id = self.language
            else:
                is_language_code = len(self.language) == 2
                raise ValueError(
                    f"Unsupported language: {self.language}. Language should be one of:"
                    f" {list(TO_LANGUAGE_CODE.values()) if is_language_code else list(TO_LANGUAGE_CODE.keys())}."
                )

        if self.task is not None:
            if self.task not in TASK_IDS:
                raise ValueError(f"Unsupported task: {self.task}. Task should be in: {TASK_IDS}")

        bos_sequence = [bos_token_id]
        if self.language is not None:
            bos_sequence.append(bos_token_id + 1 + langs.index(language_id))
        if self.task is not None:
            bos_sequence.append(transcribe_token_id if self.task == "transcribe" else translate_token_id)
        if not self.predict_timestamps:
            bos_sequence.append(notimestamps_token_id)
        return bos_sequence

    
    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

    @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 "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"

    
    def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
        self.set_prefix_tokens(task=task, language=language, predict_timestamps=not no_timestamps)
        
        
        
        
        forced_tokens = self.prefix_tokens[1:]
        forced_decoder_ids = [(rank + 1, token) for rank, token in enumerate(forced_tokens)]
        return forced_decoder_ids

    def _decode_asr(self, model_outputs, *, return_timestamps, return_language, time_precision):
        return _decode_asr(
            self,
            model_outputs,
            return_timestamps=return_timestamps,
            return_language=return_language,
            time_precision=time_precision,
        )

    
    def get_prompt_ids(self, text: str, return_tensors="np"):
        
        batch_encoding = self("<|startofprev|>", " " + text.strip(), add_special_tokens=False)

        
        prompt_text_ids = batch_encoding["input_ids"][1:]
        special_token_id = next((x for x in prompt_text_ids if x >= self.all_special_ids[0]), None)
        if special_token_id is not None:
            token = self.convert_ids_to_tokens(special_token_id)
            raise ValueError(f"Encountered text in the prompt corresponding to disallowed special token: {token}.")

        batch_encoding.convert_to_tensors(tensor_type=return_tensors)
        return batch_encoding["input_ids"]

    @staticmethod
    
    def _strip_prompt(token_ids: List[int], prompt_token_id: int, decoder_start_token_id: int):
        has_prompt = isinstance(token_ids, list) and token_ids and token_ids[0] == prompt_token_id
        if has_prompt:
            if decoder_start_token_id in token_ids:
                return token_ids[token_ids.index(decoder_start_token_id) :]
            else:
                return []

        return token_ids