codingutf8 2021 the fairseq s 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 tensorflow speech2text model from future import annotations import random from typing import optional tuple union import numpy as np import tensorflow as tf from activationstf import gettfactivation glu from modelingtfoutputs import tfbasemodeloutput tfbasemodeloutputwithpastandcrossattentions tfseq2seqlmoutput tfseq2seqmodeloutput from modelingtfutils import tfcausallanguagemodelingloss tfmodelinputtype tfpretrainedmodel tfsharedembeddings kerasserializable unpackinputs from tfutils import checkembeddingswithinbounds shapelist stablesoftmax from utils import addcodesampledocstrings addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationspeechtotext import speech2textconfig logger logging getloggername configfordoc speech2textconfig checkpointfordoc facebooks2tsmalllibrispeechasr tfspeechtotextpretrainedmodelarchivelist facebooks2tsmalllibrispeechasr see all speech2text models at https huggingface comodels filterspeechtotext largenegative 1e8 copied from transformers models bart modelingtfbart shifttokensright def shifttokensrightinputids tf tensor padtokenid int decoderstarttokenid int padtokenid tf castpadtokenid inputids dtype decoderstarttokenid tf castdecoderstarttokenid inputids dtype starttokens tf fill shapelistinputids0 1 tf converttotensordecoderstarttokenid inputids dtype shiftedinputids tf concatstarttokens inputids 1 1 replace possible 100 values in labels by padtokenid shiftedinputids tf where shiftedinputids 100 tf fillshapelistshiftedinputids tf converttotensorpadtokenid inputids dtype shiftedinputids verify that labels has only positive values and 100 assertgte0 tf debugging assertgreaterequalshiftedinputids tf constant0 dtypeinputids dtype make sure the assertion op is called by wrapping the result in an identity noop with tf controldependenciesassertgte0 shiftedinputids tf identityshiftedinputids return shiftedinputids copied from transformers models bart modelingtfbart makecausalmask def makecausalmaskinputidsshape tf tensorshape pastkeyvalueslength int 0 bsz inputidsshape0 tgtlen inputidsshape1 mask tf onestgtlen tgtlen largenegative maskcond tf rangeshapelistmask1 mask tf wheremaskcond tf reshapemaskcond 1 shapelistmask1 1 0 0 mask if pastkeyvalueslength 0 mask tf concattf zerostgtlen pastkeyvalueslength mask axis1 return tf tilemasknone none bsz 1 1 1 copied from transformers models bart modelingtfbart expandmask def expandmaskmask tf tensor tgtlen optionalint none srclen shapelistmask1 tgtlen tgtlen if tgtlen is not none else srclen onecst tf constant1 0 mask tf castmask dtypeonecst dtype expandedmask tf tilemask none none 1 1 tgtlen 1 return onecst expandedmask largenegative class tfconv1dsubsamplertf keras layers layer def initself config speech2textconfig kwargs super initkwargs self config config self numlayers config numconvlayers self inchannels config inputfeatperchannel config inputchannels self midchannels config convchannels self outchannels config dmodel self kernelsizes config convkernelsizes self convlayers tf keras layers conv1d filtersself midchannels if i self numlayers 1 else self outchannels 2 kernelsizek strides2 namefconvlayers i for i k in enumerateself kernelsizes def callself inputfeatures tf tensor tf tensor tf conv1d assumes batch x time x channels same as the input hiddenstates tf castinputfeatures tf float32 for i conv in enumerateself convlayers equivalent to paddingk 2 on pt s nn conv1d padlen self kernelsizesi 2 hiddenshapes shapelisthiddenstates hiddenstates tf concat tf zeroshiddenshapes0 padlen hiddenshapes2 hiddenstates tf zeroshiddenshapes0 padlen hiddenshapes2 axis1 hiddenstates convhiddenstates hiddenstates gluhiddenstates axis2 glu over the channel dimension return hiddenstates class tfspeech2textsinusoidalpositionalembeddingtf keras layers layer build sinusoidal embeddings this matches the implementation in tensor2tensor but differs slightly from the description in section 3 5 of attention is all you need zero pad create the position ids from the input token ids any padded tokens remain padded matt the pytorch code does a lot of work to cache the embeddings setting the cached values as a model attribute in the forward pass this is extremely forbidden in tf which wants forward calls to be idempotent tf doesn t need that caching anyway since it can just store constants during compilation so we just remove all of that code replace nonpadding symbols with their position numbers position numbers begin at paddingidx1 padding symbols are ignored this is modified from fairseq s utils makepositions args x tf tensor x returns tf tensor copied from transformers models bart modelingtfbart tfbartattention with bartspeech2text multiheaded attention from attention is all you need def init self embeddim int numheads int dropout float 0 0 isdecoder bool false bias bool true kwargs super initkwargs self embeddim embeddim self numheads numheads self dropout tf keras layers dropoutdropout self headdim embeddim numheads if self headdim numheads self embeddim raise valueerror fembeddim must be divisible by numheads got embeddim self embeddim f and numheads numheads self scaling self headdim0 5 self isdecoder isdecoder self kproj tf keras layers denseembeddim usebiasbias namekproj self qproj tf keras layers denseembeddim usebiasbias nameqproj self vproj tf keras layers denseembeddim usebiasbias namevproj self outproj tf keras layers denseembeddim usebiasbias nameoutproj def shapeself tensor tf tensor seqlen int bsz int return tf transposetf reshapetensor bsz seqlen self numheads self headdim 0 2 1 3 def call self hiddenstates tf tensor keyvaluestates tf tensor none none pastkeyvalue tupletupletf tensor none none attentionmask tf tensor none none layerheadmask tf tensor none none training optionalbool false tupletf tensor tf tensor none if keyvaluestates are provided this layer is used as a crossattention layer for the decoder get query proj get key value proj reuse k v crossattentions crossattentions reuse k v selfattention selfattention if crossattention save tupletf tensor tf tensor of all cross attention keyvaluestates further calls to crossattention layer can then reuse all crossattention keyvaluestates first if case if unidirectional selfattention decoder save tupletf tensor tf tensor of all previous decoder keyvaluestates further calls to unidirectional selfattention can concat previous decoder keyvaluestates to current projected keyvaluestates third elif case if encoder bidirectional selfattention pastkeyvalue is always none args hiddenstates tf tensor input to the layer of shape batch seqlen embeddim attentionmask tf tensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values layerheadmask tf tensor mask for attention heads in a given layer of size encoderattentionheads args hiddenstates tf tensor input to the layer of shape batch seqlen embeddim attentionmask tf tensor attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values encoderhiddenstates tf tensor cross attention input to the layer of shape batch seqlen embeddim encoderattentionmask tf tensor encoder attention mask of size batch 1 tgtlen srclen where padding elements are indicated by very large negative values layerheadmask tf tensor mask for attention heads in a given layer of size decoderattentionheads crossattnlayerheadmask tf tensor mask for heads of the crossattention module decoderattentionheads pastkeyvalue tupletf tensor cached past key and value projection states self attention decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 add present selfattn cache to positions 1 2 of presentkeyvalue tuple crossattention block crossattn cached keyvalues tuple is at positions 3 4 of presentkeyvalue tuple add crossattn to positions 3 4 of presentkeyvalue tuple fully connected computes the output length of the convolutional layers speechtotextinputsdocstring r args inputfeatures tf tensor of shape batchsize sequencelength featuresize float values of fbank features extracted from the raw speech waveform raw speech waveform can be obtained by loading a flac or wav audio file into an array of type listfloat or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into inputfeatures the autofeatureextractor should be used for extracting the fbank features padding and conversion into a tensor of floats see speech2textfeatureextractor call attentionmask tf tensor of shape 0 optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask decoderinputids tf tensor of shape batchsize targetsequencelength optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using speech2texttokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are decoder input ids glossarydecoderinputids speechtotext uses the eostokenid as the starting token for decoderinputids generation if pastkeyvalues is used optionally only the last decoderinputids have to be input see pastkeyvalues for translation and summarization training decoderinputids should be provided if no decoderinputids is provided the model will create this tensor by shifting the inputids to the right for denoising pretraining following the paper decoderattentionmask tf tensor of shape batchsize targetsequencelength optional will be made by default and ignore pad tokens it is not recommended to set this for most use cases headmask tf tensor of shape encoderlayers encoderattentionheads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoderheadmask tf tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked crossattnheadmask tf tensor of shape decoderlayers decoderattentionheads optional mask to nullify selected heads of the crossattention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoderoutputs tf floattensor optional hidden states at the output of the last layer of the encoder used in the crossattention of the decoder of shape batchsize sequencelength hiddensize is a sequence of pastkeyvalues tupletupletf tensor of length config nlayers contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if pastkeyvalues are used the user can optionally input only the last decoderinputids those that don t have their past key value states given to this model of shape batchsize 1 instead of all decoderinputids of shape batchsize sequencelength decoderinputsembeds tf floattensor of shape batchsize targetsequencelength hiddensize optional optionally instead of passing decoderinputids you can choose to directly pass an embedded representation if pastkeyvalues is used optionally only the last decoderinputsembeds have to be input see pastkeyvalues this is useful if you want more control over how to convert decoderinputids indices into associated vectors than the model s internal embedding lookup matrix usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation transformer encoder consisting of config encoderlayers self attention layers each layer is a tfspeech2textencoderlayer args config speech2textconfig computes the output length of the convolutional layers generate creates 3d attention mask because of the shape of inputfeatures convert it to 2d if thats the case args inputfeatures tf tensor of shape batchsize sequencelength featuresize float values of fbank features extracted from the raw speech waveform raw speech waveform can be obtained by loading a flac or wav audio file into an array of type listfloat or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into inputfeatures the autofeatureextractor should be used for extracting the fbank features padding and conversion into a tensor of floats see speech2textfeatureextractor call attentionmask tf tensor of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask headmask tf tensor of shape encoderlayers encoderattentionheads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple subsample attention mask if necessary check attention mask and invert bsz seqlen bsz 1 tgtseqlen srcseqlen check if headmask has a correct number of layers specified if desired add layerdrop see https arxiv orgabs1909 11556 for description transformer decoder consisting of config decoderlayers layers each layer is a tfspeech2textdecoderlayer args config speech2textconfig if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both decoderinputids and decoderinputsembeds at the same time elif inputids is not none inputshape shapelistinputids elif inputsembeds is not none inputshape shapelistinputsembeds 1 else raise valueerroryou have to specify either decoderinputids or decoderinputsembeds pastkeyvalueslength pastkeyvalueslength shapelistpastkeyvalues002 if pastkeyvalues is not none else 0 if inputsembeds is none checkembeddingswithinboundsinputids self embedtokens vocabsize inputsembeds self embedtokensinputids self embedscale else inputsembeds inputsembeds bsz seqlen bsz 1 tgtseqlen srcseqlen if inputshape1 1 combinedattentionmask makecausalmaskinputshape pastkeyvalueslengthpastkeyvalueslength else combinedattentionmask expandmask tf onesinputshape0 inputshape1 pastkeyvalueslength tgtleninputshape1 if attentionmask is not none combinedattentionmask combinedattentionmask expandmaskattentionmask tgtleninputshape1 expand encoder attention mask if encoderhiddenstates is not none and encoderattentionmask is not none bsz seqlen bsz 1 tgtseqlen srcseqlen encoderattentionmask expandmaskencoderattentionmask tgtleninputshape1 embed positions positions self embedpositionsinputids pastkeyvalueslengthpastkeyvalueslength hiddenstates inputsembeds positions hiddenstates self dropouthiddenstates trainingtraining decoder layers allhiddenstates if outputhiddenstates else none allselfattns if outputattentions else none allcrossattns if outputattentions and encoderhiddenstates is not none else none nextdecodercache if usecache else none check if headmask and crossattnheadmask have a correct number of layers specified if desired for attnmaskname attnmask in headmask headmask crossattnheadmask crossattnheadmask if attnmask is not none tf debugging assertequal shapelistattnmask0 lenself layers message fthe attnmaskname should be specified for lenself layers layers but it is for f shapelistattnmask0 for idx decoderlayer in enumerateself layers add layerdrop see https arxiv orgabs1909 11556 for description if outputhiddenstates allhiddenstates hiddenstates dropoutprobability random uniform0 1 if training and dropoutprobability self layerdrop continue pastkeyvalue pastkeyvaluesidx if pastkeyvalues is not none else none crossattnlayerheadmask crossattnheadmaskidx if crossattnheadmask is not none else none hiddenstates layerselfattn layercrossattn presentkeyvalue decoderlayer hiddenstates attentionmaskcombinedattentionmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask layerheadmaskheadmaskidx if headmask is not none else none crossattnlayerheadmaskcrossattnlayerheadmask pastkeyvaluepastkeyvalue if usecache nextdecodercache presentkeyvalue if outputattentions allselfattns layerselfattn if encoderhiddenstates is not none allcrossattns layercrossattn hiddenstates self layernormhiddenstates if outputhiddenstates allhiddenstates hiddenstates nextcache nextdecodercache if usecache else none if not returndict return hiddenstates nextcache allhiddenstates allselfattns allcrossattns else return tfbasemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluesnextcache hiddenstatesallhiddenstates attentionsallselfattns crossattentionsallcrossattns kerasserializable class tfspeech2textmainlayertf keras layers layer configclass speech2textconfig def initself config speech2textconfig kwargs super initkwargs self config config self encoder tfspeech2textencoderconfig nameencoder self decoder tfspeech2textdecoderconfig namedecoder def getinputembeddingsself return self decoder embedtokens def setinputembeddingsself newembeddings self decoder embedtokens newembeddings unpackinputs def call self inputfeaturesnone attentionmasknone decoderinputidsnone decoderattentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone encoderoutputsnone pastkeyvaluesnone decoderinputsembedsnone usecachenone outputattentionsnone outputhiddenstatesnone returndictnone trainingfalse kwargs outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict if encoderoutputs is none encoderoutputs self encoder inputfeaturesinputfeatures attentionmaskattentionmask headmaskheadmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict trainingtraining if the user passed a tuple for encoderoutputs we wrap it in a tfbasemodeloutput when returndicttrue elif returndict and not isinstanceencoderoutputs tfbasemodeloutput encoderoutputs tfbasemodeloutput lasthiddenstateencoderoutputs0 hiddenstatesencoderoutputs1 if lenencoderoutputs 1 else none attentionsencoderoutputs2 if lenencoderoutputs 2 else none if the user passed a tfbasemodeloutput for encoderoutputs we wrap it in a tuple when returndictfalse elif not returndict and not isinstanceencoderoutputs tuple encoderoutputs encoderoutputs totuple downsample encoder attention mask if attentionmask is not none encoderattentionmask self encoder getfeaturevectorattentionmask tf shapeencoderoutputs01 attentionmask else encoderattentionmask none decoder outputs consists of decfeatures pastkeyvalue dechidden decattn decoderoutputs self decoder inputidsdecoderinputids attentionmaskdecoderattentionmask encoderhiddenstatesencoderoutputs0 encoderattentionmaskencoderattentionmask headmaskdecoderheadmask crossattnheadmaskcrossattnheadmask pastkeyvaluespastkeyvalues inputsembedsdecoderinputsembeds usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict trainingtraining if not returndict return decoderoutputs encoderoutputs return tfseq2seqmodeloutput lasthiddenstatedecoderoutputs lasthiddenstate pastkeyvaluesdecoderoutputs pastkeyvalues decoderhiddenstatesdecoderoutputs hiddenstates decoderattentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions encoderlasthiddenstateencoderoutputs lasthiddenstate encoderhiddenstatesencoderoutputs hiddenstates encoderattentionsencoderoutputs attentions addstartdocstrings the bare speech2text model outputting raw hiddenstates without any specific head on top speechtotextstartdocstring class tfspeech2textmodeltfspeech2textpretrainedmodel def initself config speech2textconfig inputs kwargs super initconfig inputs kwargs self model tfspeech2textmainlayerconfig namemodel def getencoderself return self model encoder def getdecoderself return self model decoder unpackinputs addstartdocstringstomodelforwardspeechtotextinputsdocstring addcodesampledocstrings checkpointcheckpointfordoc outputtypetfseq2seqmodeloutput configclassconfigfordoc def call self inputfeatures tfmodelinputtype none none attentionmask np ndarray tf tensor none none decoderinputids np ndarray tf tensor none none decoderattentionmask np ndarray tf tensor none none headmask np ndarray tf tensor none none decoderheadmask np ndarray tf tensor none none crossattnheadmask np ndarray tf tensor none none encoderoutputs np ndarray tf tensor none none pastkeyvalues optionaltupletupleunionnp ndarray tf tensor none decoderinputsembeds np ndarray tf tensor none none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none training bool false kwargs uniontuple tfseq2seqmodeloutput outputs self model inputfeaturesinputfeatures attentionmaskattentionmask decoderinputidsdecoderinputids decoderattentionmaskdecoderattentionmask headmaskheadmask decoderheadmaskdecoderheadmask crossattnheadmaskcrossattnheadmask encoderoutputsencoderoutputs pastkeyvaluespastkeyvalues decoderinputsembedsdecoderinputsembeds usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict trainingtraining return outputs def servingoutputself output pkv tf tupleoutput pastkeyvalues1 if self config usecache else none dechs tf converttotensoroutput decoderhiddenstates if self config outputhiddenstates else none decattns tf converttotensoroutput decoderattentions if self config outputattentions else none crossattns tf converttotensoroutput crossattentions if self config outputattentions else none enchs tf converttotensoroutput encoderhiddenstates if self config outputhiddenstates else none encattns tf converttotensoroutput encoderattentions if self config outputattentions else none return tfseq2seqmodeloutput lasthiddenstateoutput lasthiddenstate pastkeyvaluespkv decoderhiddenstatesdechs decoderattentionsdecattns crossattentionscrossattns encoderlasthiddenstateoutput encoderlasthiddenstate encoderhiddenstatesenchs encoderattentionsencattns addstartdocstrings the speech2text model with a language modeling head can be used for summarization speechtotextstartdocstring class tfspeech2textforconditionalgenerationtfspeech2textpretrainedmodel tfcausallanguagemodelingloss def initself config speech2textconfig super initconfig self model tfspeech2textmainlayerconfig namemodel self lmhead tf keras layers denseself config vocabsize usebiasfalse namelmhead todo joao investigate why speech2text has numerical issues in xla generate self supportsxlageneration false def getencoderself return self model encoder def getdecoderself return self model decoder def resizetokenembeddingsself newnumtokens int tf variable newembeddings super resizetokenembeddingsnewnumtokens return newembeddings def getoutputembeddingsself return self lmhead def setoutputembeddingsself newembeddings self lmhead newembeddings unpackinputs addstartdocstringstomodelforwardspeechtotextinputsdocstring replacereturndocstringsoutputtypetfseq2seqlmoutput configclassconfigfordoc def call self inputfeatures tfmodelinputtype none none attentionmask np ndarray tf tensor none none decoderinputids np ndarray tf tensor none none decoderattentionmask np ndarray tf tensor none none headmask np ndarray tf tensor none none decoderheadmask np ndarray tf tensor none none crossattnheadmask np ndarray tf tensor none none encoderoutputs np ndarray tf tensor none none pastkeyvalues optionaltupletupleunionnp ndarray tf tensor none decoderinputsembeds np ndarray tf tensor none none labels np ndarray tf tensor none none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none training optionalbool false kwargs uniontuple tfseq2seqlmoutput r labels tf tensor of shape batchsize sequencelength optional labels for computing the masked language modeling loss indices should either be in 0 config vocabsize or 100 see inputids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocabsize returns example python import tensorflow as tf from transformers import speech2textprocessor tfspeech2textforconditionalgeneration from datasets import loaddataset import soundfile as sf model tfspeech2textforconditionalgeneration frompretrained facebooks2tsmalllibrispeechasr frompttrue processor speech2textprocessor frompretrainedfacebooks2tsmalllibrispeechasr def maptoarraybatch speech sf readbatchfile batchspeech speech return batch ds loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation ds ds mapmaptoarray ds setformattypetf inputfeatures processor dsspeech0 samplingrate16000 returntensorstf inputfeatures batch size 1 generatedids model generateinputfeatures transcription processor batchdecodegeneratedids returndict returndict if returndict is not none else self config usereturndict if labels is not none if decoderinputids is none and decoderinputsembeds is none decoderinputids shifttokensright labels self config padtokenid self config decoderstarttokenid outputs self model inputfeaturesinputfeatures attentionmaskattentionmask decoderinputidsdecoderinputids encoderoutputsencoderoutputs decoderattentionmaskdecoderattentionmask headmaskheadmask decoderheadmaskdecoderheadmask crossattnheadmaskcrossattnheadmask pastkeyvaluespastkeyvalues decoderinputsembedsdecoderinputsembeds usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict trainingtraining lmlogits self lmheadoutputs0 maskedlmloss none if labels is none else self hfcomputelosslabels lmlogits if not returndict output lmlogits outputs1 return maskedlmloss output if maskedlmloss is not none else output return tfseq2seqlmoutput lossmaskedlmloss logitslmlogits pastkeyvaluesoutputs pastkeyvalues decoderhiddenstatesoutputs decoderhiddenstates decoderattentionsoutputs decoderattentions crossattentionsoutputs crossattentions encoderlasthiddenstateoutputs encoderlasthiddenstate encoderhiddenstatesoutputs encoderhiddenstates encoderattentionsoutputs encoderattentions def servingoutputself output pkv tf tupleoutput pastkeyvalues1 if self config usecache else none dechs tf converttotensoroutput decoderhiddenstates if self config outputhiddenstates else none decattns tf converttotensoroutput decoderattentions if self config outputattentions else none crossattns tf converttotensoroutput crossattentions if self config outputattentions else none enchs tf converttotensoroutput encoderhiddenstates if self config outputhiddenstates else none encattns tf converttotensoroutput encoderattentions if self config outputattentions else none return tfseq2seqlmoutput logitsoutput logits pastkeyvaluespkv decoderhiddenstatesdechs decoderattentionsdecattns crossattentionscrossattns encoderlasthiddenstateoutput encoderlasthiddenstate encoderhiddenstatesenchs encoderattentionsencattns def prepareinputsforgeneration self decoderinputids pastkeyvaluesnone attentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone usecachenone encoderoutputsnone kwargs cut decoderinputids if past is used if pastkeyvalues is not none decoderinputids decoderinputids 1 return inputfeatures none needs to be passed to make keras layer call happy encoderoutputs encoderoutputs pastkeyvalues pastkeyvalues decoderinputids decoderinputids attentionmask attentionmask headmask headmask decoderheadmask decoderheadmask crossattnheadmask crossattnheadmask usecache usecache change this to avoid caching presumably for debugging coding utf 8 2021 the fairseq s 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 tensorflow speech2text model see all speech2text models at https huggingface co models filter speech_to_text copied from transformers models bart modeling_tf_bart shift_tokens_right replace possible 100 values in labels by pad_token_id verify that labels has only positive values and 100 make sure the assertion op is called by wrapping the result in an identity no op copied from transformers models bart modeling_tf_bart _make_causal_mask make causal mask used for bi directional self attention copied from transformers models bart modeling_tf_bart _expand_mask expands attention_mask from bsz seq_len to bsz 1 tgt_seq_len src_seq_len convolutional subsampler a stack of 1d convolution along temporal dimension followed by non linear activation via gated linear units https arxiv org abs 1911 08460 tf conv1d assumes batch x time x channels same as the input equivalent to padding k 2 on pt s nn conv1d glu over the channel dimension this module produces sinusoidal positional embeddings of any length build sinusoidal embeddings this matches the implementation in tensor2tensor but differs slightly from the description in section 3 5 of attention is all you need zero pad create the position ids from the input token ids any padded tokens remain padded matt the pytorch code does a lot of work to cache the embeddings setting the cached values as a model attribute in the forward pass this is extremely forbidden in tf which wants forward calls to be idempotent tf doesn t need that caching anyway since it can just store constants during compilation so we just remove all of that code replace non padding symbols with their position numbers position numbers begin at padding_idx 1 padding symbols are ignored this is modified from fairseq s utils make_positions args x tf tensor x returns tf tensor copied from transformers models bart modeling_tf_bart tfbartattention with bart speech2text multi headed attention from attention is all you need input shape batch x time x channel if key_value_states are provided this layer is used as a cross attention layer for the decoder get query proj get key value proj reuse k v cross_attentions cross_attentions reuse k v self_attention self_attention if cross_attention save tuple tf tensor tf tensor of all cross attention key value_states further calls to cross_attention layer can then reuse all cross attention key value_states first if case if uni directional self attention decoder save tuple tf tensor tf tensor of all previous decoder key value_states further calls to uni directional self attention can concat previous decoder key value_states to current projected key value_states third elif case if encoder bi directional self attention past_key_value is always none args hidden_states tf tensor input to the layer of shape batch seq_len embed_dim attention_mask tf tensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values layer_head_mask tf tensor mask for attention heads in a given layer of size encoder_attention_heads args hidden_states tf tensor input to the layer of shape batch seq_len embed_dim attention_mask tf tensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values encoder_hidden_states tf tensor cross attention input to the layer of shape batch seq_len embed_dim encoder_attention_mask tf tensor encoder attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values layer_head_mask tf tensor mask for attention heads in a given layer of size decoder_attention_heads cross_attn_layer_head_mask tf tensor mask for heads of the cross attention module decoder_attention_heads past_key_value tuple tf tensor cached past key and value projection states self attention decoder uni directional self attention cached key values tuple is at positions 1 2 add present self attn cache to positions 1 2 of present_key_value tuple cross attention block cross_attn cached key values tuple is at positions 3 4 of present_key_value tuple add cross attn to positions 3 4 of present_key_value tuple fully connected computes the output length of the convolutional layers this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tf keras model https www tensorflow org api_docs python tf keras model subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior tip tensorflow models and layers in transformers accept two formats as input having all inputs as keyword arguments like pytorch models or having all inputs as a list tuple or dict in the first positional argument the reason the second format is supported is that keras methods prefer this format when passing inputs to models and layers because of this support when using methods like model fit things should just work for you just pass your inputs and labels in any format that model fit supports if however you want to use the second format outside of keras methods like fit and predict such as when creating your own layers or models with the keras functional api there are three possibilities you can use to gather all the input tensors in the first positional argument a single tensor with input_ids only and nothing else model input_ids a list of varying length with one or several input tensors in the order given in the docstring model input_ids attention_mask or model input_ids attention_mask token_type_ids a dictionary with one or several input tensors associated to the input names given in the docstring model input_ids input_ids token_type_ids token_type_ids note that when creating models and layers with subclassing https keras io guides making_new_layers_and_models_via_subclassing then you don t need to worry about any of this as you can just pass inputs like you would to any other python function tip parameters config speech2textconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the tfpretrainedmodel from_pretrained method to load the model weights args input_features tf tensor of shape batch_size sequence_length feature_size float values of fbank features extracted from the raw speech waveform raw speech waveform can be obtained by loading a flac or wav audio file into an array of type list float or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into input_features the autofeatureextractor should be used for extracting the fbank features padding and conversion into a tensor of floats see speech2textfeatureextractor __call__ attention_mask tf tensor of shape 0 optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask decoder_input_ids tf tensor of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using speech2texttokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids speechtotext uses the eos_token_id as the starting token for decoder_input_ids generation if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values for translation and summarization training decoder_input_ids should be provided if no decoder_input_ids is provided the model will create this tensor by shifting the input_ids to the right for denoising pre training following the paper decoder_attention_mask tf tensor of shape batch_size target_sequence_length optional will be made by default and ignore pad tokens it is not recommended to set this for most use cases head_mask tf tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask tf tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask tf tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tf floattensor optional hidden states at the output of the last layer of the encoder used in the cross attention of the decoder of shape batch_size sequence_length hidden_size is a sequence of past_key_values tuple tuple tf tensor of length config n_layers contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length decoder_inputs_embeds tf floattensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation if past_key_values is used optionally only the last decoder_inputs_embeds have to be input see past_key_values this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation transformer encoder consisting of config encoder_layers self attention layers each layer is a tfspeech2textencoderlayer args config speech2textconfig computes the output length of the convolutional layers generate creates 3d attention mask because of the shape of input_features convert it to 2d if thats the case args input_features tf tensor of shape batch_size sequence_length feature_size float values of fbank features extracted from the raw speech waveform raw speech waveform can be obtained by loading a flac or wav audio file into an array of type list float or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into input_features the autofeatureextractor should be used for extracting the fbank features padding and conversion into a tensor of floats see speech2textfeatureextractor __call__ attention_mask tf tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask head_mask tf tensor of shape encoder_layers encoder_attention_heads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple subsample attention mask if necessary check attention mask and invert bsz seq_len bsz 1 tgt_seq_len src_seq_len check if head_mask has a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description skip the layer transformer decoder consisting of config decoder_layers layers each layer is a tfspeech2textdecoderlayer args config speech2textconfig args input_ids tf tensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary padding will be ignored by default should you provide it indices can be obtained using speech2texttokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are input ids glossary input ids attention_mask tf tensor of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask encoder_hidden_states tf tensor of shape batch_size encoder_sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask tf tensor of shape batch_size encoder_sequence_length optional mask to avoid performing cross attention on padding tokens indices of encoder input_ids mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask head_mask tf tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask tf tensor of shape decoder_layers decoder_attention_heads optional mask to nullify selected heads of the cross attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked past_key_values tuple tuple tf tensor of length config n_layers with each tuple having 2 tuples each of which has 2 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention blocks can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length inputs_embeds tf tensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional whether or not to return a utils modeloutput instead of a plain tuple past_key_values_length bsz seq_len bsz 1 tgt_seq_len src_seq_len expand encoder attention mask bsz seq_len bsz 1 tgt_seq_len src_seq_len embed positions decoder layers check if head_mask and cross_attn_head_mask have a correct number of layers specified if desired add layerdrop see https arxiv org abs 1909 11556 for description if the user passed a tuple for encoder_outputs we wrap it in a tfbasemodeloutput when return_dict true if the user passed a tfbasemodeloutput for encoder_outputs we wrap it in a tuple when return_dict false downsample encoder attention mask decoder outputs consists of dec_features past_key_value dec_hidden dec_attn todo joao investigate why speech2text has numerical issues in xla generate labels tf tensor of shape batch_size sequence_length optional labels for computing the masked language modeling loss indices should either be in 0 config vocab_size or 100 see input_ids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocab_size returns example python import tensorflow as tf from transformers import speech2textprocessor tfspeech2textforconditionalgeneration from datasets import load_dataset import soundfile as sf model tfspeech2textforconditionalgeneration from_pretrained facebook s2t small librispeech asr from_pt true processor speech2textprocessor from_pretrained facebook s2t small librispeech asr def map_to_array batch speech _ sf read batch file batch speech speech return batch ds load_dataset hf internal testing librispeech_asr_dummy clean split validation ds ds map map_to_array ds set_format type tf input_features processor ds speech 0 sampling_rate 16000 return_tensors tf input_features batch size 1 generated_ids model generate input_features transcription processor batch_decode generated_ids cut decoder_input_ids if past is used needs to be passed to make keras layer __call__ happy change this to avoid caching presumably for debugging
from __future__ import annotations

import random
from typing import Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ...activations_tf import get_tf_activation, glu
from ...modeling_tf_outputs import (
    TFBaseModelOutput,
    TFBaseModelOutputWithPastAndCrossAttentions,
    TFSeq2SeqLMOutput,
    TFSeq2SeqModelOutput,
)
from ...modeling_tf_utils import (
    TFCausalLanguageModelingLoss,
    TFModelInputType,
    TFPreTrainedModel,
    TFSharedEmbeddings,
    keras_serializable,
    unpack_inputs,
)
from ...tf_utils import check_embeddings_within_bounds, shape_list, stable_softmax
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from .configuration_speech_to_text import Speech2TextConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "Speech2TextConfig"
_CHECKPOINT_FOR_DOC = "facebook/s2t-small-librispeech-asr"


TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "facebook/s2t-small-librispeech-asr",
    
]


LARGE_NEGATIVE = -1e8



def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):
    pad_token_id = tf.cast(pad_token_id, input_ids.dtype)
    decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)
    start_tokens = tf.fill(
        (shape_list(input_ids)[0], 1), tf.convert_to_tensor(decoder_start_token_id, input_ids.dtype)
    )
    shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)
    
    shifted_input_ids = tf.where(
        shifted_input_ids == -100,
        tf.fill(shape_list(shifted_input_ids), tf.convert_to_tensor(pad_token_id, input_ids.dtype)),
        shifted_input_ids,
    )

    
    assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))

    
    with tf.control_dependencies([assert_gte0]):
        shifted_input_ids = tf.identity(shifted_input_ids)

    return shifted_input_ids



def _make_causal_mask(input_ids_shape: tf.TensorShape, past_key_values_length: int = 0):
    
    bsz = input_ids_shape[0]
    tgt_len = input_ids_shape[1]
    mask = tf.ones((tgt_len, tgt_len)) * LARGE_NEGATIVE
    mask_cond = tf.range(shape_list(mask)[-1])

    mask = tf.where(mask_cond < tf.reshape(mask_cond + 1, (shape_list(mask)[-1], 1)), 0.0, mask)

    if past_key_values_length > 0:
        mask = tf.concat([tf.zeros((tgt_len, past_key_values_length)), mask], axis=-1)

    return tf.tile(mask[None, None, :, :], (bsz, 1, 1, 1))



def _expand_mask(mask: tf.Tensor, tgt_len: Optional[int] = None):
    
    src_len = shape_list(mask)[1]
    tgt_len = tgt_len if tgt_len is not None else src_len
    one_cst = tf.constant(1.0)
    mask = tf.cast(mask, dtype=one_cst.dtype)
    expanded_mask = tf.tile(mask[:, None, None, :], (1, 1, tgt_len, 1))

    return (one_cst - expanded_mask) * LARGE_NEGATIVE


class TFConv1dSubsampler(tf.keras.layers.Layer):
    

    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.num_layers = config.num_conv_layers
        self.in_channels = config.input_feat_per_channel * config.input_channels
        self.mid_channels = config.conv_channels
        self.out_channels = config.d_model
        self.kernel_sizes = config.conv_kernel_sizes

        self.conv_layers = [
            tf.keras.layers.Conv1D(
                filters=self.mid_channels if i < self.num_layers - 1 else self.out_channels * 2,
                kernel_size=k,
                strides=2,
                name=f"conv_layers.{i}",
            )
            for i, k in enumerate(self.kernel_sizes)
        ]

    def call(self, input_features: tf.Tensor) -> tf.Tensor:
        
        hidden_states = tf.cast(input_features, tf.float32)
        for i, conv in enumerate(self.conv_layers):
            
            pad_len = self.kernel_sizes[i] // 2
            hidden_shapes = shape_list(hidden_states)
            hidden_states = tf.concat(
                (
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                    hidden_states,
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                ),
                axis=1,
            )

            hidden_states = conv(hidden_states)
            hidden_states = glu(hidden_states, axis=2)  
        return hidden_states


class TFSpeech2TextSinusoidalPositionalEmbedding(tf.keras.layers.Layer):
    

    def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None, **kwargs):
        super().__init__(**kwargs)
        self.offset = 2
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.embedding_weights = self._get_embedding(num_positions + self.offset, embedding_dim, padding_idx)

    @staticmethod
    def _get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None) -> tf.Tensor:
        
        half_dim = embedding_dim // 2
        emb = tf.math.log(10000.0) / (half_dim - 1)
        emb = tf.math.exp(tf.range(half_dim, dtype=tf.float32) * -emb)
        emb = tf.expand_dims(tf.range(num_embeddings, dtype=tf.float32), axis=1) * tf.expand_dims(emb, axis=0)
        emb = tf.reshape(tf.concat([tf.math.sin(emb), tf.math.cos(emb)], axis=1), shape=[num_embeddings, -1])
        if embedding_dim % 2 == 1:
            
            emb = tf.concat([emb, tf.zeros(num_embeddings, 1)], axis=1)
        if padding_idx is not None:
            emb = tf.concat([emb[:padding_idx, :], tf.zeros((1, tf.shape(emb)[1])), emb[padding_idx + 1 :, :]], axis=0)
        return emb

    def call(self, input_ids: tf.Tensor, past_key_values_length: int = 0) -> tf.Tensor:
        bsz, seq_len = shape_list(input_ids)
        
        position_ids = self.create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)

        
        
        
        
        embeddings = self._get_embedding(
            self.padding_idx + 1 + seq_len + self.offset + past_key_values_length, self.embedding_dim, self.padding_idx
        )
        return tf.reshape(tf.gather(embeddings, tf.reshape(position_ids, (-1,)), axis=0), (bsz, seq_len, -1))

    @staticmethod
    def create_position_ids_from_input_ids(
        input_ids: tf.Tensor, padding_idx: int, past_key_values_length: Optional[int] = 0
    ) -> tf.Tensor:
        
        mask = tf.cast(tf.math.not_equal(input_ids, padding_idx), dtype=tf.int32)
        incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask
        return tf.cast(incremental_indices, dtype=tf.int64) + padding_idx



class TFSpeech2TextAttention(tf.keras.layers.Layer):
    

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        is_decoder: bool = False,
        bias: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim

        self.num_heads = num_heads
        self.dropout = tf.keras.layers.Dropout(dropout)
        self.head_dim = embed_dim // num_heads
        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5
        self.is_decoder = is_decoder

        self.k_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="k_proj")
        self.q_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="q_proj")
        self.v_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="v_proj")
        self.out_proj = tf.keras.layers.Dense(embed_dim, use_bias=bias, name="out_proj")

    def _shape(self, tensor: tf.Tensor, seq_len: int, bsz: int):
        return tf.transpose(tf.reshape(tensor, (bsz, seq_len, self.num_heads, self.head_dim)), (0, 2, 1, 3))

    def call(
        self,
        hidden_states: tf.Tensor,
        key_value_states: tf.Tensor | None = None,
        past_key_value: Tuple[Tuple[tf.Tensor]] | None = None,
        attention_mask: tf.Tensor | None = None,
        layer_head_mask: tf.Tensor | None = None,
        training: Optional[bool] = False,
    ) -> Tuple[tf.Tensor, tf.Tensor | None]:
        

        
        
        is_cross_attention = key_value_states is not None
        bsz, tgt_len, embed_dim = shape_list(hidden_states)

        
        query_states = self.q_proj(hidden_states) * self.scaling
        
        if is_cross_attention and past_key_value is not None:
            
            key_states = past_key_value[0]
            value_states = past_key_value[1]
        elif is_cross_attention:
            
            key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
            value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
        elif past_key_value is not None:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
            key_states = tf.concat([past_key_value[0], key_states], axis=2)
            value_states = tf.concat([past_key_value[1], value_states], axis=2)
        else:
            
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        if self.is_decoder:
            
            
            
            
            
            
            
            past_key_value = (key_states, value_states)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = tf.reshape(self._shape(query_states, tgt_len, bsz), proj_shape)
        key_states = tf.reshape(key_states, proj_shape)
        value_states = tf.reshape(value_states, proj_shape)

        src_len = shape_list(key_states)[1]
        attn_weights = tf.matmul(query_states, key_states, transpose_b=True)

        tf.debugging.assert_equal(
            shape_list(attn_weights),
            [bsz * self.num_heads, tgt_len, src_len],
            message=(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {shape_list(attn_weights)}"
            ),
        )

        if attention_mask is not None:
            tf.debugging.assert_equal(
                shape_list(attention_mask),
                [bsz, 1, tgt_len, src_len],
                message=(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                    f" {shape_list(attention_mask)}"
                ),
            )

            attention_mask = tf.cast(attention_mask, dtype=attn_weights.dtype)
            attn_weights = tf.reshape(attn_weights, (bsz, self.num_heads, tgt_len, src_len)) + attention_mask
            attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len))

        attn_weights = stable_softmax(attn_weights, axis=-1)

        if layer_head_mask is not None:
            tf.debugging.assert_equal(
                shape_list(layer_head_mask),
                [self.num_heads],
                message=(
                    f"Head mask for a single layer should be of size {(self.num_heads)}, but is"
                    f" {shape_list(layer_head_mask)}"
                ),
            )

            attn_weights = tf.reshape(layer_head_mask, (1, -1, 1, 1)) * tf.reshape(
                attn_weights, (bsz, self.num_heads, tgt_len, src_len)
            )
            attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len))

        attn_probs = self.dropout(attn_weights, training=training)
        attn_output = tf.matmul(attn_probs, value_states)

        tf.debugging.assert_equal(
            shape_list(attn_output),
            [bsz * self.num_heads, tgt_len, self.head_dim],
            message=(
                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {shape_list(attn_output)}"
            ),
        )

        attn_output = tf.transpose(
            tf.reshape(attn_output, (bsz, self.num_heads, tgt_len, self.head_dim)), (0, 2, 1, 3)
        )
        attn_output = tf.reshape(attn_output, (bsz, tgt_len, embed_dim))

        attn_output = self.out_proj(attn_output)
        attn_weights: tf.Tensor = tf.reshape(attn_weights, (bsz, self.num_heads, tgt_len, src_len))

        return attn_output, attn_weights, past_key_value


class TFSpeech2TextEncoderLayer(tf.keras.layers.Layer):
    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = config.d_model
        self.self_attn = TFSpeech2TextAttention(
            self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout, name="self_attn"
        )
        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
        self.dropout = tf.keras.layers.Dropout(config.dropout)
        self.activation_fn = get_tf_activation(config.activation_function)
        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
        self.fc1 = tf.keras.layers.Dense(config.encoder_ffn_dim, name="fc1")
        self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2")
        self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm")

    def call(
        self, hidden_states: tf.Tensor, attention_mask: tf.Tensor, layer_head_mask: tf.Tensor, training: bool = False
    ):
        
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        hidden_states, self_attn_weights, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            training=training,
        )

        tf.debugging.assert_equal(
            shape_list(hidden_states),
            shape_list(residual),
            message=f"Self attn modified the shape of query {shape_list(residual)} to {shape_list(hidden_states)}",
        )

        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = self.activation_dropout(hidden_states, training=training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = residual + hidden_states

        return hidden_states, self_attn_weights


class TFSpeech2TextDecoderLayer(tf.keras.layers.Layer):
    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = config.d_model

        self.self_attn = TFSpeech2TextAttention(
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
            name="self_attn",
            is_decoder=True,
        )
        self.dropout = tf.keras.layers.Dropout(config.dropout)
        self.activation_fn = get_tf_activation(config.activation_function)
        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)

        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
        self.encoder_attn = TFSpeech2TextAttention(
            self.embed_dim,
            config.decoder_attention_heads,
            dropout=config.attention_dropout,
            name="encoder_attn",
            is_decoder=True,
        )
        self.encoder_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="encoder_attn_layer_norm")
        self.fc1 = tf.keras.layers.Dense(config.decoder_ffn_dim, name="fc1")
        self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2")
        self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm")

    def call(
        self,
        hidden_states,
        attention_mask: tf.Tensor | None = None,
        encoder_hidden_states: tf.Tensor | None = None,
        encoder_attention_mask: tf.Tensor | None = None,
        layer_head_mask: tf.Tensor | None = None,
        cross_attn_layer_head_mask: tf.Tensor | None = None,
        past_key_value: Tuple[tf.Tensor] | None = None,
        training=False,
    ) -> Tuple[tf.Tensor, tf.Tensor, Tuple[Tuple[tf.Tensor]]]:
        
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        
        hidden_states, self_attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=attention_mask,
            layer_head_mask=layer_head_mask,
            training=training,
        )
        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = residual + hidden_states

        
        cross_attn_present_key_value = None
        cross_attn_weights = None
        if encoder_hidden_states is not None:
            residual = hidden_states
            hidden_states = self.encoder_attn_layer_norm(hidden_states)

            
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
                hidden_states=hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                layer_head_mask=cross_attn_layer_head_mask,
                past_key_value=cross_attn_past_key_value,
                training=training,
            )
            hidden_states = self.dropout(hidden_states, training=training)
            hidden_states = residual + hidden_states

            
            present_key_value = present_key_value + cross_attn_present_key_value

        
        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = self.activation_dropout(hidden_states, training=training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = residual + hidden_states

        return (
            hidden_states,
            self_attn_weights,
            cross_attn_weights,
            present_key_value,
        )


class TFSpeech2TextPreTrainedModel(TFPreTrainedModel):
    config_class = Speech2TextConfig
    base_model_prefix = "model"
    main_input_name = "input_features"
    _keys_to_ignore_on_load_unexpected = [r"encoder.embed_positions.weights"]

    def _get_feat_extract_output_lengths(self, input_lengths: tf.Tensor):
        
        for _ in range(self.config.num_conv_layers):
            input_lengths = (input_lengths - 1) // 2 + 1

        return input_lengths

    @property
    def input_signature(self):
        return {
            "input_features": tf.TensorSpec(
                (None, None, self.config.input_feat_per_channel * self.config.input_channels),
                tf.float32,
                name="input_features",
            ),
            "attention_mask": tf.TensorSpec((None, None), tf.int32, name="attention_mask"),
            "decoder_input_ids": tf.TensorSpec((None, None), tf.int32, name="decoder_input_ids"),
            "decoder_attention_mask": tf.TensorSpec((None, None), tf.int32, name="decoder_attention_mask"),
        }


SPEECH_TO_TEXT_START_DOCSTRING = r


SPEECH_TO_TEXT_INPUTS_DOCSTRING = r


@keras_serializable
class TFSpeech2TextEncoder(tf.keras.layers.Layer):
    config_class = Speech2TextConfig
    

    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config

        self.dropout = tf.keras.layers.Dropout(config.dropout)
        self.layerdrop = config.encoder_layerdrop

        embed_dim = config.d_model
        self.padding_idx = config.pad_token_id
        self.max_source_positions = config.max_source_positions
        self.embed_scale = tf.math.sqrt(float(embed_dim)) if config.scale_embedding else 1.0

        self.conv = TFConv1dSubsampler(config, name="conv")

        self.embed_positions = TFSpeech2TextSinusoidalPositionalEmbedding(
            num_positions=config.max_source_positions,
            embedding_dim=embed_dim,
            padding_idx=self.padding_idx,
            name="embed_positions",
        )
        self.layers = [TFSpeech2TextEncoderLayer(config, name=f"layers.{i}") for i in range(config.encoder_layers)]
        self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm")

    def _get_feat_extract_output_lengths(self, input_lengths: tf.Tensor):
        
        for _ in range(self.config.num_conv_layers):
            input_lengths = (input_lengths - 1) // 2 + 1

        return input_lengths

    def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask):
        
        
        if len(attention_mask.shape) > 2:
            attention_mask = attention_mask[:, :, -1]

        subsampled_lengths = self._get_feat_extract_output_lengths(tf.math.reduce_sum(attention_mask, -1))
        bsz = shape_list(attention_mask)[0]
        indices = tf.concat(
            (
                tf.expand_dims(tf.range(bsz, dtype=attention_mask.dtype), -1),
                tf.expand_dims(subsampled_lengths - 1, -1),
            ),
            axis=-1,
        )
        attention_mask = tf.scatter_nd(indices=indices, updates=tf.ones(bsz), shape=[bsz, feature_vector_length])
        attention_mask = tf.cast(tf.reverse(tf.math.cumsum(tf.reverse(attention_mask, [-1]), -1), [-1]), tf.int64)
        return attention_mask

    @unpack_inputs
    def call(
        self,
        input_features=None,
        attention_mask=None,
        head_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
    ):
        
        if input_features is None:
            raise ValueError("You have to specify input_features")

        inputs_embeds = self.conv(input_features)
        inputs_embeds = self.embed_scale * inputs_embeds

        
        if attention_mask is not None:
            attention_mask = self._get_feature_vector_attention_mask(tf.shape(inputs_embeds)[1], attention_mask)
            padding_mask = tf.cast(tf.math.not_equal(attention_mask, 1), tf.int64)
        else:
            padding_mask = tf.zeros(tf.shape(inputs_embeds)[:-1], dtype=tf.int64)

        embed_pos = self.embed_positions(padding_mask)

        hidden_states = inputs_embeds + embed_pos
        hidden_states = self.dropout(hidden_states, training=training)

        
        if attention_mask is not None:
            
            attention_mask = _expand_mask(attention_mask)

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        
        if head_mask is not None:
            tf.debugging.assert_equal(
                shape_list(head_mask)[0],
                len(self.layers),
                message=(
                    f"The head_mask should be specified for {len(self.layers)} layers, but it is for"
                    f" {shape_list(head_mask)[0]}."
                ),
            )

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            
            dropout_probability = random.uniform(0, 1)
            if training and (dropout_probability < self.layerdrop):  
                continue

            hidden_states, attn = encoder_layer(
                hidden_states,
                attention_mask,
                head_mask[idx] if head_mask is not None else None,
                training=training,
            )

            if output_attentions:
                all_attentions += (attn,)

        hidden_states = self.layer_norm(hidden_states)
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        return TFBaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
        )


@keras_serializable
class TFSpeech2TextDecoder(tf.keras.layers.Layer):
    config_class = Speech2TextConfig
    

    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.layerdrop = config.decoder_layerdrop
        self.padding_idx = config.pad_token_id
        self.max_target_positions = config.max_target_positions
        self.embed_scale = tf.math.sqrt(float(config.d_model)) if config.scale_embedding else 1.0

        self.embed_tokens = TFSharedEmbeddings(config.vocab_size, config.d_model, name="embed_tokens")

        self.embed_positions = TFSpeech2TextSinusoidalPositionalEmbedding(
            num_positions=config.max_target_positions,
            embedding_dim=config.d_model,
            padding_idx=self.padding_idx,
            name="embed_positions",
        )

        self.layers = [TFSpeech2TextDecoderLayer(config, name=f"layers.{i}") for i in range(config.decoder_layers)]
        self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm")

        self.dropout = tf.keras.layers.Dropout(config.dropout)

    def get_embed_tokens(self):
        return self.embed_tokens

    def set_embed_tokens(self, embed_tokens):
        self.embed_tokens = embed_tokens

    @unpack_inputs
    def call(
        self,
        input_ids=None,
        inputs_embeds=None,
        attention_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        head_mask=None,
        cross_attn_head_mask=None,
        past_key_values=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
    ):
        r

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        
        past_key_values_length = shape_list(past_key_values[0][0])[2] if past_key_values is not None else 0

        if inputs_embeds is None:
            check_embeddings_within_bounds(input_ids, self.embed_tokens.vocab_size)
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        else:
            inputs_embeds = inputs_embeds

        
        if input_shape[-1] > 1:
            combined_attention_mask = _make_causal_mask(input_shape, past_key_values_length=past_key_values_length)
        else:
            combined_attention_mask = _expand_mask(
                tf.ones((input_shape[0], input_shape[1] + past_key_values_length)), tgt_len=input_shape[-1]
            )

        if attention_mask is not None:
            combined_attention_mask = combined_attention_mask + _expand_mask(attention_mask, tgt_len=input_shape[-1])

        
        if encoder_hidden_states is not None and encoder_attention_mask is not None:
            
            encoder_attention_mask = _expand_mask(encoder_attention_mask, tgt_len=input_shape[-1])

        
        positions = self.embed_positions(input_ids, past_key_values_length=past_key_values_length)

        hidden_states = inputs_embeds + positions
        hidden_states = self.dropout(hidden_states, training=training)

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        all_cross_attns = () if (output_attentions and encoder_hidden_states is not None) else None
        next_decoder_cache = () if use_cache else None

        
        for attn_mask_name, attn_mask in [("head_mask", head_mask), ("cross_attn_head_mask", cross_attn_head_mask)]:
            if attn_mask is not None:
                tf.debugging.assert_equal(
                    shape_list(attn_mask)[0],
                    len(self.layers),
                    message=(
                        f"The {attn_mask_name} should be specified for {len(self.layers)} layers, but it is for"
                        f" {shape_list(attn_mask)[0]}."
                    ),
                )

        for idx, decoder_layer in enumerate(self.layers):
            
            if output_hidden_states:
                all_hidden_states += (hidden_states,)
            dropout_probability = random.uniform(0, 1)
            if training and (dropout_probability < self.layerdrop):
                continue

            past_key_value = past_key_values[idx] if past_key_values is not None else None
            cross_attn_layer_head_mask = cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None

            hidden_states, layer_self_attn, layer_cross_attn, present_key_value = decoder_layer(
                hidden_states,
                attention_mask=combined_attention_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=encoder_attention_mask,
                layer_head_mask=head_mask[idx] if head_mask is not None else None,
                cross_attn_layer_head_mask=cross_attn_layer_head_mask,
                past_key_value=past_key_value,
            )

            if use_cache:
                next_decoder_cache += (present_key_value,)

            if output_attentions:
                all_self_attns += (layer_self_attn,)

                if encoder_hidden_states is not None:
                    all_cross_attns += (layer_cross_attn,)

        hidden_states = self.layer_norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None

        if not return_dict:
            return hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attns
        else:
            return TFBaseModelOutputWithPastAndCrossAttentions(
                last_hidden_state=hidden_states,
                past_key_values=next_cache,
                hidden_states=all_hidden_states,
                attentions=all_self_attns,
                cross_attentions=all_cross_attns,
            )


@keras_serializable
class TFSpeech2TextMainLayer(tf.keras.layers.Layer):
    config_class = Speech2TextConfig

    def __init__(self, config: Speech2TextConfig, **kwargs):
        super().__init__(**kwargs)
        self.config = config

        self.encoder = TFSpeech2TextEncoder(config, name="encoder")
        self.decoder = TFSpeech2TextDecoder(config, name="decoder")

    def get_input_embeddings(self):
        return self.decoder.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.decoder.embed_tokens = new_embeddings

    @unpack_inputs
    def call(
        self,
        input_features=None,
        attention_mask=None,
        decoder_input_ids=None,
        decoder_attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        encoder_outputs=None,
        past_key_values=None,
        decoder_inputs_embeds=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
        **kwargs,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_features=input_features,
                attention_mask=attention_mask,
                head_mask=head_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
                training=training,
            )
        
        elif return_dict and not isinstance(encoder_outputs, TFBaseModelOutput):
            encoder_outputs = TFBaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )
        
        elif not return_dict and not isinstance(encoder_outputs, tuple):
            encoder_outputs = encoder_outputs.to_tuple()

        
        if attention_mask is not None:
            encoder_attention_mask = self.encoder._get_feature_vector_attention_mask(
                tf.shape(encoder_outputs[0])[1], attention_mask
            )
        else:
            encoder_attention_mask = None

        
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            encoder_attention_mask=encoder_attention_mask,
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return TFSeq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )


@add_start_docstrings(
    "The bare Speech2Text Model outputting raw hidden-states without any specific head on top.",
    SPEECH_TO_TEXT_START_DOCSTRING,
)
class TFSpeech2TextModel(TFSpeech2TextPreTrainedModel):
    def __init__(self, config: Speech2TextConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.model = TFSpeech2TextMainLayer(config, name="model")

    def get_encoder(self):
        return self.model.encoder

    def get_decoder(self):
        return self.model.decoder

    @unpack_inputs
    @add_start_docstrings_to_model_forward(SPEECH_TO_TEXT_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=TFSeq2SeqModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def call(
        self,
        input_features: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        decoder_input_ids: np.ndarray | tf.Tensor | None = None,
        decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
        head_mask: np.ndarray | tf.Tensor | None = None,
        decoder_head_mask: np.ndarray | tf.Tensor | None = None,
        cross_attn_head_mask: np.ndarray | tf.Tensor | None = None,
        encoder_outputs: np.ndarray | tf.Tensor | None = None,
        past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
        decoder_inputs_embeds: np.ndarray | tf.Tensor | None = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
        **kwargs,
    ) -> Union[Tuple, TFSeq2SeqModelOutput]:
        outputs = self.model(
            input_features=input_features,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            encoder_outputs=encoder_outputs,
            past_key_values=past_key_values,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return outputs

    def serving_output(self, output):
        pkv = tf.tuple(output.past_key_values)[1] if self.config.use_cache else None
        dec_hs = tf.convert_to_tensor(output.decoder_hidden_states) if self.config.output_hidden_states else None
        dec_attns = tf.convert_to_tensor(output.decoder_attentions) if self.config.output_attentions else None
        cross_attns = tf.convert_to_tensor(output.cross_attentions) if self.config.output_attentions else None
        enc_hs = tf.convert_to_tensor(output.encoder_hidden_states) if self.config.output_hidden_states else None
        enc_attns = tf.convert_to_tensor(output.encoder_attentions) if self.config.output_attentions else None

        return TFSeq2SeqModelOutput(
            last_hidden_state=output.last_hidden_state,
            past_key_values=pkv,
            decoder_hidden_states=dec_hs,
            decoder_attentions=dec_attns,
            cross_attentions=cross_attns,
            encoder_last_hidden_state=output.encoder_last_hidden_state,
            encoder_hidden_states=enc_hs,
            encoder_attentions=enc_attns,
        )


@add_start_docstrings(
    "The Speech2Text Model with a language modeling head. Can be used for summarization.",
    SPEECH_TO_TEXT_START_DOCSTRING,
)
class TFSpeech2TextForConditionalGeneration(TFSpeech2TextPreTrainedModel, TFCausalLanguageModelingLoss):
    def __init__(self, config: Speech2TextConfig):
        super().__init__(config)
        self.model = TFSpeech2TextMainLayer(config, name="model")
        self.lm_head = tf.keras.layers.Dense(self.config.vocab_size, use_bias=False, name="lm_head")
        
        self.supports_xla_generation = False

    def get_encoder(self):
        return self.model.encoder

    def get_decoder(self):
        return self.model.decoder

    def resize_token_embeddings(self, new_num_tokens: int) -> tf.Variable:
        new_embeddings = super().resize_token_embeddings(new_num_tokens)
        return new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    @unpack_inputs
    @add_start_docstrings_to_model_forward(SPEECH_TO_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        input_features: TFModelInputType | None = None,
        attention_mask: np.ndarray | tf.Tensor | None = None,
        decoder_input_ids: np.ndarray | tf.Tensor | None = None,
        decoder_attention_mask: np.ndarray | tf.Tensor | None = None,
        head_mask: np.ndarray | tf.Tensor | None = None,
        decoder_head_mask: np.ndarray | tf.Tensor | None = None,
        cross_attn_head_mask: np.ndarray | tf.Tensor | None = None,
        encoder_outputs: np.ndarray | tf.Tensor | None = None,
        past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None,
        decoder_inputs_embeds: np.ndarray | tf.Tensor | None = None,
        labels: np.ndarray | tf.Tensor | None = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: Optional[bool] = False,
        **kwargs,
    ) -> Union[Tuple, TFSeq2SeqLMOutput]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            if decoder_input_ids is None and decoder_inputs_embeds is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        outputs = self.model(
            input_features=input_features,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )
        lm_logits = self.lm_head(outputs[0])
        masked_lm_loss = None if labels is None else self.hf_compute_loss(labels, lm_logits)

        if not return_dict:
            output = (lm_logits,) + outputs[1:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return TFSeq2SeqLMOutput(
            loss=masked_lm_loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )

    def serving_output(self, output):
        pkv = tf.tuple(output.past_key_values)[1] if self.config.use_cache else None
        dec_hs = tf.convert_to_tensor(output.decoder_hidden_states) if self.config.output_hidden_states else None
        dec_attns = tf.convert_to_tensor(output.decoder_attentions) if self.config.output_attentions else None
        cross_attns = tf.convert_to_tensor(output.cross_attentions) if self.config.output_attentions else None
        enc_hs = tf.convert_to_tensor(output.encoder_hidden_states) if self.config.output_hidden_states else None
        enc_attns = tf.convert_to_tensor(output.encoder_attentions) if self.config.output_attentions else None

        return TFSeq2SeqLMOutput(
            logits=output.logits,
            past_key_values=pkv,
            decoder_hidden_states=dec_hs,
            decoder_attentions=dec_attns,
            cross_attentions=cross_attns,
            encoder_last_hidden_state=output.encoder_last_hidden_state,
            encoder_hidden_states=enc_hs,
            encoder_attentions=enc_attns,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        past_key_values=None,
        attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        **kwargs,
    ):
        
        if past_key_values is not None:
            decoder_input_ids = decoder_input_ids[:, -1:]

        return {
            "input_features": None,  
            "encoder_outputs": encoder_outputs,
            "past_key_values": past_key_values,
            "decoder_input_ids": decoder_input_ids,
            "attention_mask": attention_mask,
            "head_mask": head_mask,
            "decoder_head_mask": decoder_head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "use_cache": use_cache,  
        }