codingutf8 2021 the facebook inc 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 tokenization class for wav2vec2 import json import os import sys import warnings from dataclasses import dataclass from itertools import groupby from typing import typechecking dict list optional tuple union import numpy as np from tokenizationutils import pretrainedtokenizer from tokenizationutilsbase import addedtoken batchencoding from utils import modeloutput paddingstrategy tensortype addenddocstrings isflaxavailable istfavailable istorchavailable logging topyobj logger logging getloggername if typechecking if istorchavailable import torch if istfavailable import tensorflow as tf if isflaxavailable import jax numpy as jnp noqa f401 vocabfilesnames vocabfile vocab json tokenizerconfigfile tokenizerconfig json pretrainedvocabfilesmap vocabfile facebookwav2vec2base960h https huggingface cofacebookwav2vec2base960hresolvemainvocab json tokenizerconfigfile facebookwav2vec2base960h https huggingface cofacebookwav2vec2base960hresolvemaintokenizerconfig json wav2vec2 has no max input length pretrainedpositionalembeddingssizes facebookwav2vec2base960h sys maxsize wav2vec2kwargsdocstring r padding bool str or utils paddingstrategy optional defaults to false activates and controls padding accepts the following values true or longest pad to the longest sequence in the batch or no padding if only a single sequence if provided maxlength pad to a maximum length specified with the argument maxlength or to the maximum acceptable input length for the model if that argument is not provided false or donotpad default no padding i e can output a batch with sequences of different lengths maxlength int optional controls the maximum length to use by one of the truncationpadding parameters if left unset or set to none this will use the predefined model maximum length if a maximum length is required by one of the truncationpadding parameters if the model has no specific maximum input length like xlnet truncationpadding to a maximum length will be deactivated padtomultipleof int optional if set will pad the sequence to a multiple of the provided value this is especially useful to enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta returntensors str or utils tensortype optional if set will return tensors instead of list of python integers acceptable values are tf return tensorflow tf constant objects pt return pytorch torch tensor objects np return numpy np ndarray objects verbose bool optional defaults to true whether or not to print more information and warnings output type of wav2vec2ctctokenizer with transcription args text list of str or str decoded logits in text from usually the speech transcription charoffsets list of listdictstr unionint str or listdictstr unionint str offsets of the decoded characters in combination with sampling rate and model downsampling rate char offsets can be used to compute time stamps for each charater total logit score of the beam associated with produced text wordoffsets list of listdictstr unionint str or listdictstr unionint str offsets of the decoded words in combination with sampling rate and model downsampling rate word offsets can be used to compute time stamps for each word constructs a wav2vec2ctc 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 vocabfile str file containing the vocabulary bostoken str optional defaults to s the beginning of sentence token eostoken str optional defaults to s the end of sentence token unktoken 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 padtoken str optional defaults to pad the token used for padding for example when batching sequences of different lengths worddelimitertoken str optional defaults to the token used for defining the end of a word dolowercase bool optional defaults to false whether or not to accept lowercase input and lowercase the output when decoding targetlang str optional a target language the tokenizer should set by default targetlang has to be defined for multilingual nested vocabulary such as facebookmms1ballhttps huggingface cofacebookmms1ball kwargs additional keyword arguments passed along to pretrainedtokenizer if target lang is defined vocab must be a nested dict with each target lang being one vocabulary make sure that tokens made of several characters are not split at tokenization set the target language of a nested multilingual dictionary make sure that tokens made of several characters are not split at tokenization str word delimiter token log an error if used while not having been set optionalint id of the worddelimitertoken in the vocabulary returns none if the token has not been set overwritten to never strip converts a string in a sequence of tokens string using the tokenizer converts a token str in an index integer using the vocab return self encoder gettoken self encoder getself unktoken def convertidtotokenself index int str converts a connectionisttemporalclassification ctc output tokens into a single string group same tokens into nonrepeating tokens in ctc style decoding filter self padtoken which is used as ctcblank token replace delimiter token retrieve offsets set tokens to correct processed token retrieve word offsets from character offsets don t output chars if not set to true join to string filter out ctc token if we are in the same state as before we simply repeat what we ve done before switching state finishing a word starting a new word special decode function is needed for wav2vec2tokenizer because added tokens should be treated exactly the same as tokens of the base vocabulary and therefore the function converttokenstostring has to be called on the whole token list and not individually on added tokens overwritten from tokenizationutilsbase py because tokenizer can output modeloutput which should not be a list for batched output and because we need docs for outputcharoffsets here convert a list of lists of token ids into a list of strings by calling decode args sequences unionlistint listlistint np ndarray torch tensor tf tensor list of tokenized input ids can be obtained using the call method skipspecialtokens bool optional defaults to false whether or not to remove special tokens in the decoding cleanuptokenizationspaces bool optional whether or not to clean up the tokenization spaces outputcharoffsets bool optional defaults to false whether or not to output character offsets character offsets can be used in combination with the sampling rate and model downsampling rate to compute the timestamps of transcribed characters tip please take a look at the example of wav2vec2ctctokenizer decode to better understand how to make use of outputcharoffsets wav2vec2ctctokenizer batchdecode works the same way with batched output tip outputwordoffsets bool optional defaults to false whether or not to output word offsets word offsets can be used in combination with the sampling rate and model downsampling rate to compute the timestamps of transcribed words tip please take a look at the example of wav2vec2ctctokenizer decode to better understand how to make use of outputwordoffsets wav2vec2ctctokenizer batchdecode works the same way with batched output tip kwargs additional keyword arguments optional will be passed to the underlying model specific decode method returns liststr or models wav2vec2 tokenizationwav2vec2 wav2vec2ctctokenizeroutput the list of decoded sentences will be a models wav2vec2 tokenizationwav2vec2 wav2vec2ctctokenizeroutput when outputcharoffsets true or outputwordoffsets true transform list of dicts to dict of lists overwritten from tokenizationutilsbase py because we need docs for outputcharoffsets and outputwordoffsets here 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 converttokenstostringself convertidstotokenstokenids args tokenids unionint listint np ndarray torch tensor tf tensor list of tokenized input ids can be obtained using the call method skipspecialtokens bool optional defaults to false whether or not to remove special tokens in the decoding cleanuptokenizationspaces bool optional whether or not to clean up the tokenization spaces outputcharoffsets bool optional defaults to false whether or not to output character offsets character offsets can be used in combination with the sampling rate and model downsampling rate to compute the timestamps of transcribed characters tip please take a look at the example below to better understand how to make use of outputcharoffsets tip outputwordoffsets bool optional defaults to false whether or not to output word offsets word offsets can be used in combination with the sampling rate and model downsampling rate to compute the timestamps of transcribed words tip please take a look at the example below to better understand how to make use of outputwordoffsets tip kwargs additional keyword arguments optional will be passed to the underlying model specific decode method returns str or models wav2vec2 tokenizationwav2vec2 wav2vec2ctctokenizeroutput the list of decoded sentences will be a models wav2vec2 tokenizationwav2vec2 wav2vec2ctctokenizeroutput when outputcharoffsets true or outputwordoffsets true example python let s see how to retrieve time steps for a model from transformers import autotokenizer autofeatureextractor automodelforctc from datasets import loaddataset import datasets import torch import model feature extractor tokenizer model automodelforctc frompretrainedfacebookwav2vec2base960h tokenizer autotokenizer frompretrainedfacebookwav2vec2base960h featureextractor autofeatureextractor frompretrainedfacebookwav2vec2base960h load first sample of english commonvoice dataset loaddatasetmozillafoundationcommonvoice110 en splittrain streamingtrue dataset dataset castcolumnaudio datasets audiosamplingrate16000 datasetiter iterdataset sample nextdatasetiter forward sample through model to get greedily predicted transcription ids inputvalues featureextractorsampleaudioarray returntensorspt inputvalues logits modelinputvalues logits0 predids torch argmaxlogits axis1 retrieve word stamps analogous commands for outputcharoffsets outputs tokenizer decodepredids outputwordoffsetstrue compute timeoffset in seconds as product of downsampling ratio and samplingrate timeoffset model config inputstologitsratio featureextractor samplingrate wordoffsets word dword starttime rounddstartoffset timeoffset 2 endtime rounddendoffset timeoffset 2 for d in outputs wordoffsets compare word offsets with audio entrain0commonvoiceen19121553 mp3 online on the dataset viewer https huggingface codatasetsmozillafoundationcommonvoice110vieweren wordoffsets 3 word the starttime 0 7 endtime 0 78 word trick starttime 0 88 endtime 1 08 word appears starttime 1 2 endtime 1 64 convert inputs to python lists tokenids topyobjtokenids return self decode tokenidstokenids skipspecialtokensskipspecialtokens cleanuptokenizationspacescleanuptokenizationspaces outputcharoffsetsoutputcharoffsets outputwordoffsetsoutputwordoffsets kwargs def savevocabularyself savedirectory str filenameprefix optionalstr none tuplestr if not os path isdirsavedirectory logger errorfvocabulary path savedirectory should be a directory return vocabfile os path join savedirectory filenameprefix if filenameprefix else vocabfilesnamesvocabfile with openvocabfile w encodingutf8 as f f writejson dumpsself vocab indent2 sortkeystrue ensureasciifalse n return vocabfile class wav2vec2tokenizerpretrainedtokenizer vocabfilesnames vocabfilesnames pretrainedvocabfilesmap vocabfile facebookwav2vec2base960h https huggingface cofacebookwav2vec2base960hresolvemainvocab json tokenizerconfigfile facebookwav2vec2base960h https huggingface cofacebookwav2vec2base960hresolvemaintokenizer json modelinputnames inputvalues attentionmask def init self vocabfile bostokens eostokens unktokenunk padtokenpad worddelimitertoken dolowercasefalse donormalizefalse returnattentionmaskfalse kwargs warnings warn the class wav2vec2tokenizer is deprecated and will be removed in version 5 of transformers please use wav2vec2processor or wav2vec2ctctokenizer instead futurewarning self worddelimitertoken worddelimitertoken self dolowercase dolowercase self returnattentionmask returnattentionmask self donormalize donormalize with openvocabfile encodingutf8 as vocabhandle self encoder json loadvocabhandle self decoder v k for k v in self encoder items super init unktokenunktoken bostokenbostoken eostokeneostoken padtokenpadtoken dolowercasedolowercase donormalizedonormalize returnattentionmaskreturnattentionmask worddelimitertokenworddelimitertoken kwargs property def worddelimitertokenself str if self worddelimitertoken is none and self verbose logger errorusing worddelimitertoken but it is not set yet return none return strself worddelimitertoken property def worddelimitertokenidself optionalint if self worddelimitertoken is none return none return self converttokenstoidsself worddelimitertoken worddelimitertoken setter def worddelimitertokenself value self worddelimitertoken value worddelimitertokenid setter def worddelimitertokenidself value self worddelimitertoken self converttokenstoidsvalue addenddocstringswav2vec2kwargsdocstring def call self rawspeech unionnp ndarray listfloat listnp ndarray listlistfloat padding unionbool str paddingstrategy false maxlength optionalint none padtomultipleof optionalint none returntensors optionalunionstr tensortype none verbose bool true kwargs batchencoding isbatchednumpy isinstancerawspeech np ndarray and lenrawspeech shape 1 if isbatchednumpy and lenrawspeech shape 2 raise valueerrorfonly monochannel audio is supported for input to self isbatched isbatchednumpy or isinstancerawspeech list tuple and isinstancerawspeech0 np ndarray tuple list make sure input is in list format if isbatched and not isinstancerawspeech0 np ndarray rawspeech np asarrayspeech for speech in rawspeech elif not isbatched and not isinstancerawspeech np ndarray rawspeech np asarrayrawspeech always return batch if not isbatched rawspeech rawspeech zeromean and unitvariance normalization if self donormalize rawspeech x np meanx np sqrtnp varx 1e5 for x in rawspeech convert into correct format for padding encodedinputs batchencodinginputvalues rawspeech paddedinputs self pad encodedinputs paddingpadding maxlengthmaxlength padtomultipleofpadtomultipleof returnattentionmaskself returnattentionmask returntensorsreturntensors verboseverbose return paddedinputs property def vocabsizeself int return lenself decoder def getvocabself dict return dictself encoder self addedtokensencoder def converttokentoidself token str int converts an index integer in a token str using the vocab result self decoder getindex self unktoken return result def converttokenstostringself tokens liststr str group same tokens into nonrepeating tokens in ctc style decoding groupedtokens tokengroup0 for tokengroup in groupbytokens filter self padtoken which is used as ctcblank token filteredtokens listfilterlambda token token self padtoken groupedtokens replace delimiter token string join if token self worddelimitertoken else token for token in filteredtokens strip if self dolowercase string string lower return string def decode self tokenids listint skipspecialtokens bool false cleanuptokenizationspaces bool none kwargs str filteredtokens self convertidstotokenstokenids skipspecialtokensskipspecialtokens result for token in filteredtokens if skipspecialtokens and token in self allspecialids continue result appendtoken text self converttokenstostringresult cleanuptokenizationspaces cleanuptokenizationspaces if cleanuptokenizationspaces is not none else self cleanuptokenizationspaces if cleanuptokenizationspaces cleantext self cleanuptokenizationtext return cleantext else return text def savevocabularyself savedirectory str filenameprefix optionalstr none tuplestr if not os path isdirsavedirectory logger errorfvocabulary path savedirectory should be a directory return vocabfile os path join savedirectory filenameprefix if filenameprefix else vocabfilesnamesvocabfile with openvocabfile w encodingutf8 as f f writejson dumpsself encoder indent2 sortkeystrue ensureasciifalse n return vocabfile coding utf 8 2021 the facebook inc 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 tokenization class for wav2vec2 noqa f401 wav2vec2 has no max input length padding bool str or utils paddingstrategy optional defaults to false activates and controls padding accepts the following values true or longest pad to the longest sequence in the batch or no padding if only a single sequence if provided max_length pad to a maximum length specified with the argument max_length or to the maximum acceptable input length for the model if that argument is not provided false or do_not_pad default no padding i e can output a batch with sequences of different lengths max_length int optional controls the maximum length to use by one of the truncation padding parameters if left unset or set to none this will use the predefined model maximum length if a maximum length is required by one of the truncation padding parameters if the model has no specific maximum input length like xlnet truncation padding to a maximum length will be deactivated pad_to_multiple_of int optional if set will pad the sequence to a multiple of the provided value this is especially useful to enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta return_tensors str or utils tensortype optional if set will return tensors instead of list of python integers acceptable values are tf return tensorflow tf constant objects pt return pytorch torch tensor objects np return numpy np ndarray objects verbose bool optional defaults to true whether or not to print more information and warnings output type of wav2vec2ctctokenizer with transcription args text list of str or str decoded logits in text from usually the speech transcription char_offsets list of list dict str union int str or list dict str union int str offsets of the decoded characters in combination with sampling rate and model downsampling rate char offsets can be used to compute time stamps for each charater total logit score of the beam associated with produced text word_offsets list of list dict str union int str or list dict str union int str offsets of the decoded words in combination with sampling rate and model downsampling rate word offsets can be used to compute time stamps for each word constructs a wav2vec2ctc 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 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 word_delimiter_token str optional defaults to the token used for defining the end of a word do_lower_case bool optional defaults to false whether or not to accept lowercase input and lowercase the output when decoding target_lang str optional a target language the tokenizer should set by default target_lang has to be defined for multi lingual nested vocabulary such as facebook mms 1b all https huggingface co facebook mms 1b all kwargs additional keyword arguments passed along to pretrainedtokenizer if target lang is defined vocab must be a nested dict with each target lang being one vocabulary make sure that tokens made of several characters are not split at tokenization set the target language of a nested multi lingual dictionary make sure that tokens made of several characters are not split at tokenization str word delimiter token log an error if used while not having been set optional int id of the word_delimiter_token in the vocabulary returns none if the token has not been set overwritten to never strip converts a string in a sequence of tokens string using the tokenizer converts a token str in an index integer using the vocab converts an index integer in a token str using the vocab converts a connectionist temporal classification ctc output tokens into a single string group same tokens into non repeating tokens in ctc style decoding filter self pad_token which is used as ctc blank token replace delimiter token retrieve offsets set tokens to correct processed token retrieve word offsets from character offsets don t output chars if not set to true join to string filter out ctc token if we are in the same state as before we simply repeat what we ve done before switching state finishing a word starting a new word special _decode function is needed for wav2vec2tokenizer because added tokens should be treated exactly the same as tokens of the base vocabulary and therefore the function convert_tokens_to_string has to be called on the whole token list and not individually on added tokens overwritten from tokenization_utils_base py because tokenizer can output modeloutput which should not be a list for batched output and because we need docs for output_char_offsets here convert a list of lists of token ids into a list of strings by calling decode args sequences union list int list 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 output_char_offsets bool optional defaults to false whether or not to output character offsets character offsets can be used in combination with the sampling rate and model downsampling rate to compute the time stamps of transcribed characters tip please take a look at the example of wav2vec2ctctokenizer decode to better understand how to make use of output_char_offsets wav2vec2ctctokenizer batch_decode works the same way with batched output tip output_word_offsets bool optional defaults to false whether or not to output word offsets word offsets can be used in combination with the sampling rate and model downsampling rate to compute the time stamps of transcribed words tip please take a look at the example of wav2vec2ctctokenizer decode to better understand how to make use of output_word_offsets wav2vec2ctctokenizer batch_decode works the same way with batched output tip kwargs additional keyword arguments optional will be passed to the underlying model specific decode method returns list str or models wav2vec2 tokenization_wav2vec2 wav2vec2ctctokenizeroutput the list of decoded sentences will be a models wav2vec2 tokenization_wav2vec2 wav2vec2ctctokenizeroutput when output_char_offsets true or output_word_offsets true transform list of dicts to dict of lists overwritten from tokenization_utils_base py because we need docs for output_char_offsets and output_word_offsets here 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 output_char_offsets bool optional defaults to false whether or not to output character offsets character offsets can be used in combination with the sampling rate and model downsampling rate to compute the time stamps of transcribed characters tip please take a look at the example below to better understand how to make use of output_char_offsets tip output_word_offsets bool optional defaults to false whether or not to output word offsets word offsets can be used in combination with the sampling rate and model downsampling rate to compute the time stamps of transcribed words tip please take a look at the example below to better understand how to make use of output_word_offsets tip kwargs additional keyword arguments optional will be passed to the underlying model specific decode method returns str or models wav2vec2 tokenization_wav2vec2 wav2vec2ctctokenizeroutput the list of decoded sentences will be a models wav2vec2 tokenization_wav2vec2 wav2vec2ctctokenizeroutput when output_char_offsets true or output_word_offsets true example python let s see how to retrieve time steps for a model from transformers import autotokenizer autofeatureextractor automodelforctc from datasets import load_dataset import datasets import torch import model feature extractor tokenizer model automodelforctc from_pretrained facebook wav2vec2 base 960h tokenizer autotokenizer from_pretrained facebook wav2vec2 base 960h feature_extractor autofeatureextractor from_pretrained facebook wav2vec2 base 960h load first sample of english common_voice dataset load_dataset mozilla foundation common_voice_11_0 en split train streaming true dataset dataset cast_column audio datasets audio sampling_rate 16_000 dataset_iter iter dataset sample next dataset_iter forward sample through model to get greedily predicted transcription ids input_values feature_extractor sample audio array return_tensors pt input_values logits model input_values logits 0 pred_ids torch argmax logits axis 1 retrieve word stamps analogous commands for output_char_offsets outputs tokenizer decode pred_ids output_word_offsets true compute time_offset in seconds as product of downsampling ratio and sampling_rate time_offset model config inputs_to_logits_ratio feature_extractor sampling_rate word_offsets word d word start_time round d start_offset time_offset 2 end_time round d end_offset time_offset 2 for d in outputs word_offsets compare word offsets with audio en_train_0 common_voice_en_19121553 mp3 online on the dataset viewer https huggingface co datasets mozilla foundation common_voice_11_0 viewer en word_offsets 3 word the start_time 0 7 end_time 0 78 word trick start_time 0 88 end_time 1 08 word appears start_time 1 2 end_time 1 64 convert inputs to python lists constructs a wav2vec2 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 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 word_delimiter_token str optional defaults to the token used for defining the end of a word do_lower_case bool optional defaults to false whether or not to lowercase the output when decoding do_normalize bool optional defaults to false whether or not to zero mean unit variance normalize the input normalizing can help to significantly improve the performance for some models e g wav2vec2 lv60 https huggingface co models search lv60 return_attention_mask bool optional defaults to false whether or not wav2vec2tokenizer __call__ should return attention_mask tip wav2vec2 models that have set config feat_extract_norm group such as wav2vec2 base https huggingface co facebook wav2vec2 base 960h have not been trained using attention_mask for such models input_values should simply be padded with 0 and no attention_mask should be passed for wav2vec2 models that have set config feat_extract_norm layer such as wav2vec2 lv60 https huggingface co facebook wav2vec2 large 960h lv60 self attention_mask should be passed for batched inference tip kwargs additional keyword arguments passed along to pretrainedtokenizer str padding token log an error if used while not having been set optional int id of the word_delimiter_token in the vocabulary returns none if the token has not been set main method to tokenize and prepare for the model one or several sequence s or one or several pair s of sequences args raw_speech np ndarray list float list np ndarray list list float the sequence or batch of sequences to be padded each sequence can be a numpy array a list of float values a list of numpy array or a list of list of float values must be mono channel audio not stereo i e single float per timestep make sure input is in list format always return batch zero mean and unit variance normalization convert into correct format for padding converts a token str in an index integer using the vocab converts an index integer in a token str using the vocab converts a connectionist temporal classification ctc output tokens into a single string group same tokens into non repeating tokens in ctc style decoding filter self pad_token which is used as ctc blank token replace delimiter token special _decode function is needed for wav2vec2tokenizer because added tokens should be treated exactly the same as tokens of the base vocabulary and therefore the function convert_tokens_to_string has to be called on the whole token list and not individually on added tokens
import json
import os
import sys
import warnings
from dataclasses import dataclass
from itertools import groupby
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import numpy as np

from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...utils import (
    ModelOutput,
    PaddingStrategy,
    TensorType,
    add_end_docstrings,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    logging,
    to_py_obj,
)


logger = logging.get_logger(__name__)


if TYPE_CHECKING:
    if is_torch_available():
        import torch
    if is_tf_available():
        import tensorflow as tf
    if is_flax_available():
        import jax.numpy as jnp  


VOCAB_FILES_NAMES = {
    "vocab_file": "vocab.json",
    "tokenizer_config_file": "tokenizer_config.json",
}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {
        "facebook/wav2vec2-base-960h": "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/vocab.json",
    },
    "tokenizer_config_file": {
        "facebook/wav2vec2-base-960h": (
            "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/tokenizer_config.json"
        ),
    },
}


PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {"facebook/wav2vec2-base-960h": sys.maxsize}

WAV2VEC2_KWARGS_DOCSTRING = r

ListOfDict = List[Dict[str, Union[int, str]]]


@dataclass
class Wav2Vec2CTCTokenizerOutput(ModelOutput):
    

    text: Union[List[str], str]
    char_offsets: Union[List[ListOfDict], ListOfDict] = None
    word_offsets: Union[List[ListOfDict], ListOfDict] = None


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

    def __init__(
        self,
        vocab_file,
        bos_token="<s>",
        eos_token="</s>",
        unk_token="<unk>",
        pad_token="<pad>",
        word_delimiter_token="|",
        replace_word_delimiter_char=" ",
        do_lower_case=False,
        target_lang=None,
        **kwargs,
    ):
        self._word_delimiter_token = word_delimiter_token

        self.do_lower_case = do_lower_case
        self.replace_word_delimiter_char = replace_word_delimiter_char
        self.target_lang = target_lang

        with open(vocab_file, encoding="utf-8") as vocab_handle:
            self.vocab = json.load(vocab_handle)

        
        
        if target_lang is not None:
            self.encoder = self.vocab[target_lang]
        else:
            self.encoder = self.vocab

        self.decoder = {v: k for k, v in self.encoder.items()}

        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            do_lower_case=do_lower_case,
            word_delimiter_token=word_delimiter_token,
            replace_word_delimiter_char=replace_word_delimiter_char,
            target_lang=target_lang,
            **kwargs,
        )

        
        
        for token in self.encoder.keys():
            if len(token) > 1:
                self.add_tokens(AddedToken(token, rstrip=True, lstrip=True, normalized=False))

    def set_target_lang(self, target_lang: str):
        
        if self.vocab == self.encoder:
            raise ValueError(f"{self.vocab} is not a multi-lingual, nested tokenizer. Cannot set target language.")

        if target_lang not in self.vocab:
            raise ValueError(f"{target_lang} does not exist. Choose one of {', '.join(self.vocab.keys())}.")

        self.target_lang = target_lang
        self.init_kwargs["target_lang"] = target_lang
        self.encoder = self.vocab[target_lang]
        self.decoder = {v: k for k, v in self.encoder.items()}

        
        
        for token in self.encoder.keys():
            if len(token) > 1:
                self.add_tokens(AddedToken(token, rstrip=True, lstrip=True, normalized=False))

    @property
    def word_delimiter_token(self) -> str:
        
        if self._word_delimiter_token is None and self.verbose:
            logger.error("Using word_delimiter_token, but it is not set yet.")
            return None
        return str(self._word_delimiter_token)

    @property
    def word_delimiter_token_id(self) -> Optional[int]:
        
        if self._word_delimiter_token is None:
            return None
        return self.convert_tokens_to_ids(self.word_delimiter_token)

    @word_delimiter_token.setter
    def word_delimiter_token(self, value):
        self._word_delimiter_token = value

    @word_delimiter_token_id.setter
    def word_delimiter_token_id(self, value):
        self._word_delimiter_token = self.convert_tokens_to_ids(value)

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

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

    def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
        
        to_add = []
        for token in new_tokens:
            if isinstance(token, str):
                to_add.append(AddedToken(token, rstrip=False, lstrip=False, normalize=False))
            else:
                to_add.append(token)

        return super()._add_tokens(to_add, special_tokens)

    def _tokenize(self, text, **kwargs):
        
        if self.do_lower_case:
            text = text.upper()

        return list(text.replace(" ", self.word_delimiter_token))

    def _convert_token_to_id(self, token: str) -> int:
        
        return self.encoder.get(token, self.encoder.get(self.unk_token))

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

    def convert_tokens_to_string(
        self,
        tokens: List[str],
        group_tokens: bool = True,
        spaces_between_special_tokens: bool = False,
        output_char_offsets: bool = False,
        output_word_offsets: bool = False,
    ) -> Dict[str, Union[str, float]]:
        
        if len(tokens) == 0:
            return {"text": "", "char_offsets": [], "word_offsets": []}
        
        if group_tokens:
            chars, char_repetitions = zip(*((token, len(list(group_iter))) for token, group_iter in groupby(tokens)))
        else:
            chars = tokens
            char_repetitions = len(tokens) * [1]

        
        processed_chars = list(filter(lambda char: char != self.pad_token, chars))

        
        processed_chars = [
            self.replace_word_delimiter_char if char == self.word_delimiter_token else char for char in processed_chars
        ]

        
        char_offsets = word_offsets = None
        if output_char_offsets or output_word_offsets:
            char_offsets = self._compute_offsets(char_repetitions, chars, self.pad_token)

            if len(char_offsets) != len(processed_chars):
                raise ValueError(
                    f"`char_offsets`: {char_offsets} and `processed_tokens`: {processed_chars}"
                    " have to be of the same length, but are: "
                    f"`len(offsets)`: {len(char_offsets)} and `len(processed_tokens)`:"
                    f" {len(processed_chars)}"
                )

            
            for i, char in enumerate(processed_chars):
                char_offsets[i]["char"] = char

            
            word_offsets = None
            if output_word_offsets:
                word_offsets = self._get_word_offsets(char_offsets, self.replace_word_delimiter_char)

            
            if not output_char_offsets:
                char_offsets = None

        
        join_char = " " if spaces_between_special_tokens else ""
        string = join_char.join(processed_chars).strip()

        if self.do_lower_case:
            string = string.lower()

        return {"text": string, "char_offsets": char_offsets, "word_offsets": word_offsets}

    @staticmethod
    def _compute_offsets(
        char_repetitions: List[int], chars: List[str], ctc_token: int
    ) -> List[Dict[str, Union[str, int]]]:
        end_indices = np.asarray(char_repetitions).cumsum()
        start_indices = np.concatenate(([0], end_indices[:-1]))

        offsets = [
            {"char": t, "start_offset": s, "end_offset": e} for t, s, e in zip(chars, start_indices, end_indices)
        ]

        
        offsets = list(filter(lambda offsets: offsets["char"] != ctc_token, offsets))
        return offsets

    @staticmethod
    def _get_word_offsets(
        offsets: Dict[str, Union[str, float]], word_delimiter_char: str = " "
    ) -> Dict[str, Union[str, float]]:
        word_offsets = []

        last_state = "SPACE"
        word = ""
        start_offset = 0
        end_offset = 0
        for i, offset in enumerate(offsets):
            char = offset["char"]
            state = "SPACE" if char == word_delimiter_char else "WORD"

            if state == last_state:
                
                end_offset = offset["end_offset"]
                word += char
            else:
                
                if state == "SPACE":
                    
                    word_offsets.append({"word": word, "start_offset": start_offset, "end_offset": end_offset})
                else:
                    
                    start_offset = offset["start_offset"]
                    end_offset = offset["end_offset"]
                    word = char

            last_state = state
        if last_state == "WORD":
            word_offsets.append({"word": word, "start_offset": start_offset, "end_offset": end_offset})

        return word_offsets

    def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
        if is_split_into_words:
            text = " " + text
        return (text, kwargs)

    def _decode(
        self,
        token_ids: List[int],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        group_tokens: bool = True,
        spaces_between_special_tokens: bool = False,
        output_word_offsets: Optional[bool] = False,
        output_char_offsets: Optional[bool] = False,
    ) -> str:
        
        filtered_tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens)

        result = []
        for token in filtered_tokens:
            if skip_special_tokens and token in self.all_special_ids:
                continue
            result.append(token)

        string_output = self.convert_tokens_to_string(
            result,
            group_tokens=group_tokens,
            spaces_between_special_tokens=spaces_between_special_tokens,
            output_word_offsets=output_word_offsets,
            output_char_offsets=output_char_offsets,
        )

        text = string_output["text"]

        clean_up_tokenization_spaces = (
            clean_up_tokenization_spaces
            if clean_up_tokenization_spaces is not None
            else self.clean_up_tokenization_spaces
        )
        if clean_up_tokenization_spaces:
            text = self.clean_up_tokenization(text)

        if output_word_offsets or output_char_offsets:
            return Wav2Vec2CTCTokenizerOutput(
                text=text,
                char_offsets=string_output["char_offsets"],
                word_offsets=string_output["word_offsets"],
            )
        else:
            return text

    
    
    
    def batch_decode(
        self,
        sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        output_char_offsets: bool = False,
        output_word_offsets: bool = False,
        **kwargs,
    ) -> List[str]:
        
        batch_decoded = [
            self.decode(
                seq,
                skip_special_tokens=skip_special_tokens,
                clean_up_tokenization_spaces=clean_up_tokenization_spaces,
                output_char_offsets=output_char_offsets,
                output_word_offsets=output_word_offsets,
                **kwargs,
            )
            for seq in sequences
        ]
        if output_char_offsets or output_word_offsets:
            
            return Wav2Vec2CTCTokenizerOutput({k: [d[k] for d in batch_decoded] for k in batch_decoded[0]})

        return batch_decoded

    
    
    def decode(
        self,
        token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        output_char_offsets: bool = False,
        output_word_offsets: bool = False,
        **kwargs,
    ) -> str:
        
        
        token_ids = to_py_obj(token_ids)

        return self._decode(
            token_ids=token_ids,
            skip_special_tokens=skip_special_tokens,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            output_char_offsets=output_char_offsets,
            output_word_offsets=output_word_offsets,
            **kwargs,
        )

    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
        vocab_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        )

        with open(vocab_file, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n")

        return (vocab_file,)


class Wav2Vec2Tokenizer(PreTrainedTokenizer):
    

    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = {
        "vocab_file": {
            "facebook/wav2vec2-base-960h": "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/vocab.json"
        },
        "tokenizer_config_file": {
            "facebook/wav2vec2-base-960h": (
                "https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/tokenizer.json"
            ),
        },
    }
    model_input_names = ["input_values", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        bos_token="<s>",
        eos_token="</s>",
        unk_token="<unk>",
        pad_token="<pad>",
        word_delimiter_token="|",
        do_lower_case=False,
        do_normalize=False,
        return_attention_mask=False,
        **kwargs,
    ):
        warnings.warn(
            "The class `Wav2Vec2Tokenizer` is deprecated and will be removed in version 5 of Transformers. Please use"
            " `Wav2Vec2Processor` or `Wav2Vec2CTCTokenizer` instead.",
            FutureWarning,
        )

        self._word_delimiter_token = word_delimiter_token

        self.do_lower_case = do_lower_case
        self.return_attention_mask = return_attention_mask
        self.do_normalize = do_normalize

        with open(vocab_file, encoding="utf-8") as vocab_handle:
            self.encoder = json.load(vocab_handle)

        self.decoder = {v: k for k, v in self.encoder.items()}

        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            do_lower_case=do_lower_case,
            do_normalize=do_normalize,
            return_attention_mask=return_attention_mask,
            word_delimiter_token=word_delimiter_token,
            **kwargs,
        )

    @property
    def word_delimiter_token(self) -> str:
        
        if self._word_delimiter_token is None and self.verbose:
            logger.error("Using word_delimiter_token, but it is not set yet.")
            return None
        return str(self._word_delimiter_token)

    @property
    def word_delimiter_token_id(self) -> Optional[int]:
        
        if self._word_delimiter_token is None:
            return None
        return self.convert_tokens_to_ids(self.word_delimiter_token)

    @word_delimiter_token.setter
    def word_delimiter_token(self, value):
        self._word_delimiter_token = value

    @word_delimiter_token_id.setter
    def word_delimiter_token_id(self, value):
        self._word_delimiter_token = self.convert_tokens_to_ids(value)

    @add_end_docstrings(WAV2VEC2_KWARGS_DOCSTRING)
    def __call__(
        self,
        raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
        padding: Union[bool, str, PaddingStrategy] = False,
        max_length: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        verbose: bool = True,
        **kwargs,
    ) -> BatchEncoding:
        

        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
        if is_batched_numpy and len(raw_speech.shape) > 2:
            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
        is_batched = is_batched_numpy or (
            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
        )

        
        if is_batched and not isinstance(raw_speech[0], np.ndarray):
            raw_speech = [np.asarray(speech) for speech in raw_speech]
        elif not is_batched and not isinstance(raw_speech, np.ndarray):
            raw_speech = np.asarray(raw_speech)

        
        if not is_batched:
            raw_speech = [raw_speech]

        
        if self.do_normalize:
            raw_speech = [(x - np.mean(x)) / np.sqrt(np.var(x) + 1e-5) for x in raw_speech]

        
        encoded_inputs = BatchEncoding({"input_values": raw_speech})

        padded_inputs = self.pad(
            encoded_inputs,
            padding=padding,
            max_length=max_length,
            pad_to_multiple_of=pad_to_multiple_of,
            return_attention_mask=self.return_attention_mask,
            return_tensors=return_tensors,
            verbose=verbose,
        )

        return padded_inputs

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

    def get_vocab(self) -> Dict:
        return dict(self.encoder, **self.added_tokens_encoder)

    def _convert_token_to_id(self, token: str) -> int:
        
        return self.encoder.get(token, self.encoder.get(self.unk_token))

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

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        
        
        grouped_tokens = [token_group[0] for token_group in groupby(tokens)]

        
        filtered_tokens = list(filter(lambda token: token != self.pad_token, grouped_tokens))

        
        string = "".join([" " if token == self.word_delimiter_token else token for token in filtered_tokens]).strip()

        if self.do_lower_case:
            string = string.lower()

        return string

    def _decode(
        self,
        token_ids: List[int],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        **kwargs,
    ) -> str:
        
        filtered_tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens)

        result = []
        for token in filtered_tokens:
            if skip_special_tokens and token in self.all_special_ids:
                continue
            result.append(token)

        text = self.convert_tokens_to_string(result)

        clean_up_tokenization_spaces = (
            clean_up_tokenization_spaces
            if clean_up_tokenization_spaces is not None
            else self.clean_up_tokenization_spaces
        )
        if clean_up_tokenization_spaces:
            clean_text = self.clean_up_tokenization(text)
            return clean_text
        else:
            return text

    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
        vocab_file = os.path.join(
            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        )

        with open(vocab_file, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")

        return (vocab_file,)