codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch speech2text model import copy import inspect import os import tempfile import unittest from transformers import speech2textconfig from transformers testingutils import istorchavailable requiresentencepiece requiretokenizers requiretorch requiretorchfp16 requiretorchaudio slow torchdevice from transformers utils import cachedproperty from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin configzeroinit floatstensor idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import speech2textforconditionalgeneration speech2textmodel speech2textprocessor from transformers models speechtotext modelingspeechtotext import speech2textdecoder speech2textencoder def preparespeechtotextinputsdict config inputfeatures decoderinputids attentionmasknone decoderattentionmasknone headmasknone decoderheadmasknone crossattnheadmasknone if attentionmask is none attentionmask inputfeatures ne0 if decoderattentionmask is none decoderattentionmask decoderinputids neconfig padtokenid if headmask is none headmask torch onesconfig encoderlayers config encoderattentionheads devicetorchdevice if decoderheadmask is none decoderheadmask torch onesconfig decoderlayers config decoderattentionheads devicetorchdevice if crossattnheadmask is none crossattnheadmask torch onesconfig decoderlayers config decoderattentionheads devicetorchdevice return inputids inputfeatures inputfeatures inputfeatures decoderinputids decoderinputids attentionmask attentionmask decoderattentionmask attentionmask headmask headmask decoderheadmask decoderheadmask crossattnheadmask crossattnheadmask requiretorch class speech2textmodeltester def init self parent batchsize13 seqlength7 istrainingtrue uselabelsfalse vocabsize99 hiddensize16 numhiddenlayers2 numattentionheads4 intermediatesize4 numconvlayers2 convkernelsizes5 5 convchannels32 inputfeatperchannel24 inputchannels1 hiddenactrelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 maxpositionembeddings20 maxsourcepositions20 maxtargetpositions20 eostokenid2 padtokenid1 bostokenid0 self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self numconvlayers numconvlayers self convkernelsizes convkernelsizes self convchannels convchannels self inputfeatperchannel inputfeatperchannel self inputchannels inputchannels self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self maxpositionembeddings maxpositionembeddings self maxsourcepositions maxsourcepositions self maxtargetpositions maxtargetpositions self eostokenid eostokenid self padtokenid padtokenid self bostokenid bostokenid def prepareconfigandinputsself inputfeatures floatstensor self batchsize self seqlength self inputfeatperchannel self vocabsize attentionmask torch onesself batchsize self seqlength dtypetorch long devicetorchdevice decoderinputids idstensorself batchsize self seqlength self vocabsize clamp2 config self getconfig inputsdict preparespeechtotextinputsdict config inputfeaturesinputfeatures decoderinputidsdecoderinputids attentionmaskattentionmask return config inputsdict def getconfigself return speech2textconfig vocabsizeself vocabsize dmodelself hiddensize encoderlayersself numhiddenlayers decoderlayersself numhiddenlayers encoderattentionheadsself numattentionheads decoderattentionheadsself numattentionheads encoderffndimself intermediatesize decoderffndimself intermediatesize numconvlayersself numconvlayers convkernelsizesself convkernelsizes convchannelsself convchannels inputfeatperchannelself inputfeatperchannel inputchannelsself inputchannels dropoutself hiddendropoutprob attentiondropoutself attentionprobsdropoutprob maxpositionembeddingsself maxpositionembeddings maxsourcepositionsself maxsourcepositions maxtargetpositionsself maxtargetpositions eostokenidself eostokenid bostokenidself bostokenid padtokenidself padtokenid def prepareconfigandinputsforcommonself config inputsdict self prepareconfigandinputs return config inputsdict def getsubsampledoutputlengthsself inputlengths for i in rangeself numconvlayers inputlengths inputlengths 1 2 1 return inputlengths def createandcheckmodelforwardself config inputsdict model speech2textmodelconfigconfig totorchdevice eval inputfeatures inputsdictinputfeatures decoderinputids inputsdictdecoderinputids first forward pass lasthiddenstate modelinputfeatures decoderinputidsdecoderinputids lasthiddenstate self parent asserttruelasthiddenstate shape 13 7 16 def createandcheckdecodermodelpastlargeinputsself config inputsdict model speech2textmodelconfigconfig getdecoder totorchdevice eval inputids inputsdictdecoderinputids attentionmask inputsdictdecoderattentionmask first forward pass outputs modelinputids attentionmaskattentionmask usecachetrue output pastkeyvalues outputs totuple create hypothetical multiple next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize clamp2 nextattnmask idstensorself batchsize 3 2 append to next inputids and nextinputids torch catinputids nexttokens dim1 nextattentionmask torch catattentionmask nextattnmask dim1 outputfromnopast modelnextinputids attentionmasknextattentionmasklasthiddenstate outputfrompast modelnexttokens attentionmasknextattentionmask pastkeyvaluespastkeyvalues lasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach self parent asserttrueoutputfrompastslice shape1 nexttokens shape1 test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e2 def checkencoderdecodermodelstandaloneself config inputsdict model speech2textmodelconfigconfig totorchdevice eval outputs modelinputsdict encoderlasthiddenstate outputs encoderlasthiddenstate lasthiddenstate outputs lasthiddenstate with tempfile temporarydirectory as tmpdirname encoder model getencoder encoder savepretrainedtmpdirname encoder speech2textencoder frompretrainedtmpdirname totorchdevice encoderlasthiddenstate2 encoder inputsdictinputfeatures attentionmaskinputsdictattentionmask 0 self parent asserttrueencoderlasthiddenstate2 encoderlasthiddenstate abs max item 1e3 with tempfile temporarydirectory as tmpdirname decoder model getdecoder decoder savepretrainedtmpdirname decoder speech2textdecoder frompretrainedtmpdirname totorchdevice encoderattentionmask encoder getfeaturevectorattentionmask encoderlasthiddenstate shape1 inputsdictattentionmask lasthiddenstate2 decoder inputidsinputsdictdecoderinputids attentionmaskinputsdictdecoderattentionmask encoderhiddenstatesencoderlasthiddenstate encoderattentionmaskencoderattentionmask 0 self parent asserttruelasthiddenstate2 lasthiddenstate abs max item 1e3 requiretorch class speech2textmodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses speech2textmodel speech2textforconditionalgeneration if istorchavailable else allgenerativemodelclasses speech2textforconditionalgeneration if istorchavailable else pipelinemodelmapping automaticspeechrecognition speech2textforconditionalgeneration featureextraction speech2textmodel if istorchavailable else isencoderdecoder true fxcompatible true testpruning false testmissingkeys false inputname inputfeatures def setupself self modeltester speech2textmodeltesterself self configtester configtesterself configclassspeech2textconfig self maxdiff 3000 def testconfigself self configtester runcommontests def testsaveloadstrictself config inputsdict self modeltester prepareconfigandinputs for modelclass in self allmodelclasses model modelclassconfig with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname model2 info modelclass frompretrainedtmpdirname outputloadinginfotrue self assertequalinfomissingkeys def testmodelforwardself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelforwardconfigandinputs def testdecodermodelpastwithlargeinputsself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelpastlargeinputsconfigandinputs def testencoderdecodermodelstandaloneself configandinputs self modeltester prepareconfigandinputsforcommon self modeltester checkencoderdecodermodelstandaloneconfigandinputs not implemented currently def testinputsembedsself pass training is not supported yet def testtrainingself pass def testtraininggradientcheckpointingself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantself pass unittest skip reasonthis architecure seem to not compute gradients properly when using gc check https github comhuggingfacetransformerspull27124 def testtraininggradientcheckpointingusereentrantfalseself pass requiretorchfp16 def testgeneratefp16self config inputdict self modeltester prepareconfigandinputs inputfeatures inputdictinputfeatures attentionmask inputdictattentionmask model speech2textforconditionalgenerationconfig eval totorchdevice inputfeatures inputfeatures half model half model generateinputfeatures attentionmaskattentionmask model generateinputfeatures numbeams4 dosampletrue earlystoppingfalse numreturnsequences3 def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames inputfeatures attentionmask decoderinputids decoderattentionmask expectedargnames extend headmask decoderheadmask crossattnheadmask encoderoutputs if headmask and decoderheadmask and crossattnheadmask in argnames else encoderoutputs self assertlistequalargnames lenexpectedargnames expectedargnames def testhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass hiddenstates outputs encoderhiddenstates if config isencoderdecoder else outputs hiddenstates expectednumlayers getattr self modeltester expectednumhiddenlayers self modeltester numhiddenlayers 1 self assertequallenhiddenstates expectednumlayers if hasattrself modeltester encoderseqlength seqlength self modeltester encoderseqlength else seqlength self modeltester seqlength subsampledseqlength model getfeatextractoutputlengthsseqlength self assertlistequal listhiddenstates0 shape2 subsampledseqlength self modeltester hiddensize if config isencoderdecoder hiddenstates outputs decoderhiddenstates self assertisinstancehiddenstates list tuple self assertequallenhiddenstates expectednumlayers seqlen getattrself modeltester seqlength none decoderseqlength getattrself modeltester decoderseqlength seqlen self assertlistequal listhiddenstates0 shape2 decoderseqlength self modeltester hiddensize config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true seqlen getattrself modeltester seqlength none decoderseqlength getattrself modeltester decoderseqlength seqlen encoderseqlength getattrself modeltester encoderseqlength seqlen decoderkeylength getattrself modeltester decoderkeylength decoderseqlength encoderkeylength getattrself modeltester keylength encoderseqlength for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig model totorchdevice model eval subsampledencoderseqlength model getfeatextractoutputlengthsencoderseqlength subsampledencoderkeylength model getfeatextractoutputlengthsencoderkeylength with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs encoderattentions if config isencoderdecoder else outputs attentions self assertequallenattentions self modeltester numhiddenlayers check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs encoderattentions if config isencoderdecoder else outputs attentions self assertequallenattentions self modeltester numhiddenlayers self assertlistequal listattentions0 shape3 self modeltester numattentionheads subsampledencoderseqlength subsampledencoderkeylength outlen lenoutputs correctoutlen 5 loss is at first position if labels in inputsdict correctoutlen 1 loss is added to beginning if pastkeyvalues in outputs correctoutlen 1 pastkeyvalues have been returned self assertequaloutlen correctoutlen decoder attentions decoderattentions outputs decoderattentions self assertisinstancedecoderattentions list tuple self assertequallendecoderattentions self modeltester numhiddenlayers self assertlistequal listdecoderattentions0 shape3 self modeltester numattentionheads decoderseqlength decoderkeylength cross attentions crossattentions outputs crossattentions self assertisinstancecrossattentions list tuple self assertequallencrossattentions self modeltester numhiddenlayers self assertlistequal listcrossattentions0 shape3 self modeltester numattentionheads decoderseqlength subsampledencoderkeylength check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass addedhiddenstates 2 self assertequaloutlen addedhiddenstates lenoutputs selfattentions outputs encoderattentions if config isencoderdecoder else outputs attentions self assertequallenselfattentions self modeltester numhiddenlayers self assertlistequal listselfattentions0 shape3 self modeltester numattentionheads subsampledencoderseqlength subsampledencoderkeylength def testresizetokensembeddingsself originalconfig inputsdict self modeltester prepareconfigandinputsforcommon if not self testresizeembeddings return for modelclass in self allmodelclasses config copy deepcopyoriginalconfig model modelclassconfig model totorchdevice if self modeltester istraining is false model eval modelvocabsize config vocabsize retrieve the embeddings and clone theme modelembed model resizetokenembeddingsmodelvocabsize clonedembeddings modelembed weight clone check that resizing the token embeddings with a larger vocab size increases the model s vocab size modelembed model resizetokenembeddingsmodelvocabsize 10 self assertequalmodel config vocabsize modelvocabsize 10 check that it actually resizes the embeddings matrix self assertequalmodelembed weight shape0 clonedembeddings shape0 10 check that the model can still do a forward pass successfully every parameter should be resized modelself prepareforclassinputsdict modelclass check that resizing the token embeddings with a smaller vocab size decreases the model s vocab size modelembed model resizetokenembeddingsmodelvocabsize 15 self assertequalmodel config vocabsize modelvocabsize 15 check that it actually resizes the embeddings matrix self assertequalmodelembed weight shape0 clonedembeddings shape0 15 make sure that decoderinputids are resized if decoderinputids in inputsdict inputsdictdecoderinputids clampmaxmodelvocabsize 15 1 modelself prepareforclassinputsdict modelclass check that adding and removing tokens has not modified the first part of the embedding matrix modelsequal true for p1 p2 in zipclonedembeddings modelembed weight if p1 data nep2 data sum 0 modelsequal false self asserttruemodelsequal def testresizeembeddingsuntiedself originalconfig inputsdict self modeltester prepareconfigandinputsforcommon if not self testresizeembeddings return originalconfig tiewordembeddings false if model cannot untied embeddings leave test if originalconfig tiewordembeddings return for modelclass in self allmodelclasses config copy deepcopyoriginalconfig model modelclassconfig totorchdevice if no output embeddings leave test if model getoutputembeddings is none continue check that resizing the token embeddings with a larger vocab size increases the model s vocab size modelvocabsize config vocabsize model resizetokenembeddingsmodelvocabsize 10 self assertequalmodel config vocabsize modelvocabsize 10 outputembeds model getoutputembeddings self assertequaloutputembeds weight shape0 modelvocabsize 10 check bias if present if outputembeds bias is not none self assertequaloutputembeds bias shape0 modelvocabsize 10 check that the model can still do a forward pass successfully every parameter should be resized modelself prepareforclassinputsdict modelclass check that resizing the token embeddings with a smaller vocab size decreases the model s vocab size model resizetokenembeddingsmodelvocabsize 15 self assertequalmodel config vocabsize modelvocabsize 15 check that it actually resizes the embeddings matrix outputembeds model getoutputembeddings self assertequaloutputembeds weight shape0 modelvocabsize 15 check bias if present if outputembeds bias is not none self assertequaloutputembeds bias shape0 modelvocabsize 15 check that the model can still do a forward pass successfully every parameter should be resized if decoderinputids in inputsdict inputsdictdecoderinputids clampmaxmodelvocabsize 15 1 check that the model can still do a forward pass successfully every parameter should be resized modelself prepareforclassinputsdict modelclass def testgeneratewithoutinputidsself pass staticmethod def getencoderoutputs model inputids attentionmask outputattentionsnone outputhiddenstatesnone numinterleave1 encoder model getencoder encoderoutputs encoder inputids attentionmaskattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates encoderoutputslasthiddenstate encoderoutputs lasthiddenstate repeatinterleave numinterleave dim0 inputids inputids 0 inputids torch zeroslikeinputids 1 dtypetorch long model getdecoderstarttokenid attentionmask none return encoderoutputs inputids attentionmask def checkoutputsself output inputids config usecachefalse numreturnsequences1 batchsize seqlength inputids shape 2 subsampledseqlength self modeltester getsubsampledoutputlengthsseqlength numsequencesinoutput batchsize numreturnsequences genlen output sequences shape1 1 if config isencoderdecoder else output sequences shape1 seqlength scores self checkscoresnumsequencesinoutput output scores lengthgenlen configconfig attentions encoder self checkencoderattentionforgenerate output encoderattentions batchsize config subsampledseqlength decoder self checkattentionsforgenerate numsequencesinoutput output decoderattentions minlength1 maxlengthoutput sequences shape1 configconfig usecacheusecache hidden states encoder self checkencoderhiddenstatesforgenerate output encoderhiddenstates batchsize config subsampledseqlength decoder self checkhiddenstatesforgenerate numsequencesinoutput output decoderhiddenstates minlength1 maxlengthoutput sequences shape1 configconfig usecacheusecache def createandchecktorchscriptself config inputsdict if not self testtorchscript return configsnoinit configzeroinitconfig to be sure we have no nan configsnoinit torchscript true for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit model totorchdevice model eval inputs self prepareforclassinputsdict modelclass try model config usecache false fstm still requires this hack fstm should probably be refactored similar to bart afterward inputfeatures inputsinputfeatures attentionmask inputsattentionmask decoderinputids inputsdecoderinputids decoderattentionmask inputsdecoderattentionmask tracedmodel torch jit trace model inputfeatures attentionmask decoderinputids decoderattentionmask except runtimeerror self failcouldn t trace module with tempfile temporarydirectory as tmpdirname ptfilename os path jointmpdirname tracedmodel pt try torch jit savetracedmodel ptfilename except exception self failcouldn t save module try loadedmodel torch jit loadptfilename except exception self failcouldn t load module model totorchdevice model eval loadedmodel totorchdevice loadedmodel eval modelstatedict model statedict loadedmodelstatedict loadedmodel statedict nonpersistentbuffers for key in loadedmodelstatedict keys if key not in modelstatedict keys nonpersistentbufferskey loadedmodelstatedictkey loadedmodelstatedict key value for key value in loadedmodelstatedict items if key not in nonpersistentbuffers self assertequalsetmodelstatedict keys setloadedmodelstatedict keys modelbuffers listmodel buffers for nonpersistentbuffer in nonpersistentbuffers values foundbuffer false for i modelbuffer in enumeratemodelbuffers if torch equalnonpersistentbuffer modelbuffer foundbuffer true break self asserttruefoundbuffer modelbuffers popi modelsequal true for layername p1 in modelstatedict items p2 loadedmodelstatedictlayername if p1 data nep2 data sum 0 modelsequal false self asserttruemodelsequal def testpttfmodelequivalenceself allowmissingkeystrue allow missing keys since tf doesn t cache the sinusoidal embeddings in an attribute super testpttfmodelequivalenceallowmissingkeysallowmissingkeys unittest skiptest failing rocketnight is looking into it def testtffromptsafetensorsself pass requiretorch requiretorchaudio requiresentencepiece requiretokenizers slow class speech2textmodelintegrationtestsunittest testcase cachedproperty def defaultprocessorself return speech2textprocessor frompretrainedfacebooks2tsmalllibrispeechasr def loaddatasamplesself numsamples from datasets import loaddataset ds loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation automatic decoding with librispeech speechsamples ds sortid selectrangenumsamples numsamplesaudio return xarray for x in speechsamples def testgenerationlibrispeechself model speech2textforconditionalgeneration frompretrainedfacebooks2tsmalllibrispeechasr model totorchdevice processor self defaultprocessor inputspeech self loaddatasamples1 inputfeatures processorinputspeech returntensorspt inputfeatures totorchdevice generatedids model generateinputfeatures generatedtranscript processor batchdecodegeneratedids skipspecialtokenstrue expectedtranscriptions mister quilter is the apostle of the middle classes and we are glad to welcome his gospel self assertlistequalgeneratedtranscript expectedtranscriptions def testgenerationlibrispeechbatchedself model speech2textforconditionalgeneration frompretrainedfacebooks2tsmalllibrispeechasr model totorchdevice processor self defaultprocessor inputspeech self loaddatasamples4 inputs processorinputspeech returntensorspt paddingtrue inputfeatures inputs inputfeatures totorchdevice attentionmask inputs attentionmask totorchdevice generatedids model generateinputfeatures attentionmaskattentionmask generatedtranscripts processor batchdecodegeneratedids skipspecialtokenstrue expectedtranscriptions mister quilter is the apostle of the middle classes and we are glad to welcome his gospel nor is mister cultar s manner less interesting than his matter he tells us that at this festive season of the year with christmas and roast beef looming before us similes drawn from eating and its results occur most readily to the mind he has grave doubts whether sir frederick leyton s work is really greek after all and can discover in it but little of rocky ithaca self assertlistequalgeneratedtranscripts expectedtranscriptions coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch speech2text model input_ids input_features computes the output length of the convolutional layers first forward pass first forward pass create hypothetical multiple next token and extent to next_input_ids append to next input_ids and select random slice test that outputs are equal for slice not implemented currently training is not supported yet signature parameters is an ordereddict so arg_names order is deterministic check that output_hidden_states also work using config check that output_attentions also work using config loss is at first position loss is added to beginning past_key_values have been returned decoder attentions cross attentions check attention is always last and order is fine retrieve the embeddings and clone theme check that resizing the token embeddings with a larger vocab size increases the model s vocab size check that it actually resizes the embeddings matrix check that the model can still do a forward pass successfully every parameter should be resized check that resizing the token embeddings with a smaller vocab size decreases the model s vocab size check that it actually resizes the embeddings matrix make sure that decoder_input_ids are resized check that adding and removing tokens has not modified the first part of the embedding matrix if model cannot untied embeddings leave test if no output embeddings leave test check that resizing the token embeddings with a larger vocab size increases the model s vocab size check bias if present check that the model can still do a forward pass successfully every parameter should be resized check that resizing the token embeddings with a smaller vocab size decreases the model s vocab size check that it actually resizes the embeddings matrix check bias if present check that the model can still do a forward pass successfully every parameter should be resized check that the model can still do a forward pass successfully every parameter should be resized scores attentions encoder decoder hidden states encoder decoder to be sure we have no nan fstm still requires this hack fstm should probably be refactored similar to bart afterward allow missing keys since tf doesn t cache the sinusoidal embeddings in an attribute automatic decoding with librispeech
import copy
import inspect
import os
import tempfile
import unittest

from transformers import Speech2TextConfig
from transformers.testing_utils import (
    is_torch_available,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    require_torch_fp16,
    require_torchaudio,
    slow,
    torch_device,
)
from transformers.utils import cached_property

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import Speech2TextForConditionalGeneration, Speech2TextModel, Speech2TextProcessor
    from transformers.models.speech_to_text.modeling_speech_to_text import Speech2TextDecoder, Speech2TextEncoder


def prepare_speech_to_text_inputs_dict(
    config,
    input_features,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = input_features.ne(0)
    if decoder_attention_mask is None:
        decoder_attention_mask = decoder_input_ids.ne(config.pad_token_id)
    if head_mask is None:
        head_mask = torch.ones(config.encoder_layers, config.encoder_attention_heads, device=torch_device)
    if decoder_head_mask is None:
        decoder_head_mask = torch.ones(config.decoder_layers, config.decoder_attention_heads, device=torch_device)
    if cross_attn_head_mask is None:
        cross_attn_head_mask = torch.ones(config.decoder_layers, config.decoder_attention_heads, device=torch_device)
    return {
        
        "input_features": input_features,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }


@require_torch
class Speech2TextModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=7,
        is_training=True,
        use_labels=False,
        vocab_size=99,
        hidden_size=16,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=4,
        num_conv_layers=2,
        conv_kernel_sizes=(5, 5),
        conv_channels=32,
        input_feat_per_channel=24,
        input_channels=1,
        hidden_act="relu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=20,
        max_source_positions=20,
        max_target_positions=20,
        eos_token_id=2,
        pad_token_id=1,
        bos_token_id=0,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.num_conv_layers = num_conv_layers
        self.conv_kernel_sizes = conv_kernel_sizes
        self.conv_channels = conv_channels
        self.input_feat_per_channel = input_feat_per_channel
        self.input_channels = input_channels
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.max_source_positions = max_source_positions
        self.max_target_positions = max_target_positions
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id

    def prepare_config_and_inputs(self):
        input_features = floats_tensor(
            [self.batch_size, self.seq_length, self.input_feat_per_channel], self.vocab_size
        )
        attention_mask = torch.ones([self.batch_size, self.seq_length], dtype=torch.long, device=torch_device)
        decoder_input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size).clamp(2)

        config = self.get_config()
        inputs_dict = prepare_speech_to_text_inputs_dict(
            config,
            input_features=input_features,
            decoder_input_ids=decoder_input_ids,
            attention_mask=attention_mask,
        )
        return config, inputs_dict

    def get_config(self):
        return Speech2TextConfig(
            vocab_size=self.vocab_size,
            d_model=self.hidden_size,
            encoder_layers=self.num_hidden_layers,
            decoder_layers=self.num_hidden_layers,
            encoder_attention_heads=self.num_attention_heads,
            decoder_attention_heads=self.num_attention_heads,
            encoder_ffn_dim=self.intermediate_size,
            decoder_ffn_dim=self.intermediate_size,
            num_conv_layers=self.num_conv_layers,
            conv_kernel_sizes=self.conv_kernel_sizes,
            conv_channels=self.conv_channels,
            input_feat_per_channel=self.input_feat_per_channel,
            input_channels=self.input_channels,
            dropout=self.hidden_dropout_prob,
            attention_dropout=self.attention_probs_dropout_prob,
            max_position_embeddings=self.max_position_embeddings,
            max_source_positions=self.max_source_positions,
            max_target_positions=self.max_target_positions,
            eos_token_id=self.eos_token_id,
            bos_token_id=self.bos_token_id,
            pad_token_id=self.pad_token_id,
        )

    def prepare_config_and_inputs_for_common(self):
        config, inputs_dict = self.prepare_config_and_inputs()
        return config, inputs_dict

    def get_subsampled_output_lengths(self, input_lengths):
        

        for i in range(self.num_conv_layers):
            input_lengths = (input_lengths - 1) // 2 + 1

        return input_lengths

    def create_and_check_model_forward(self, config, inputs_dict):
        model = Speech2TextModel(config=config).to(torch_device).eval()

        input_features = inputs_dict["input_features"]
        decoder_input_ids = inputs_dict["decoder_input_ids"]

        
        last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state

        self.parent.assertTrue(last_hidden_state.shape, (13, 7, 16))

    def create_and_check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = Speech2TextModel(config=config).get_decoder().to(torch_device).eval()
        input_ids = inputs_dict["decoder_input_ids"]
        attention_mask = inputs_dict["decoder_attention_mask"]

        
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size).clamp(2)
        next_attn_mask = ids_tensor((self.batch_size, 3), 2)

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        next_attention_mask = torch.cat([attention_mask, next_attn_mask], dim=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)["last_hidden_state"]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[
            "last_hidden_state"
        ]

        
        random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
        output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

        
        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-2))

    def check_encoder_decoder_model_standalone(self, config, inputs_dict):
        model = Speech2TextModel(config=config).to(torch_device).eval()
        outputs = model(**inputs_dict)

        encoder_last_hidden_state = outputs.encoder_last_hidden_state
        last_hidden_state = outputs.last_hidden_state

        with tempfile.TemporaryDirectory() as tmpdirname:
            encoder = model.get_encoder()
            encoder.save_pretrained(tmpdirname)
            encoder = Speech2TextEncoder.from_pretrained(tmpdirname).to(torch_device)

        encoder_last_hidden_state_2 = encoder(
            inputs_dict["input_features"], attention_mask=inputs_dict["attention_mask"]
        )[0]

        self.parent.assertTrue((encoder_last_hidden_state_2 - encoder_last_hidden_state).abs().max().item() < 1e-3)

        with tempfile.TemporaryDirectory() as tmpdirname:
            decoder = model.get_decoder()
            decoder.save_pretrained(tmpdirname)
            decoder = Speech2TextDecoder.from_pretrained(tmpdirname).to(torch_device)

        encoder_attention_mask = encoder._get_feature_vector_attention_mask(
            encoder_last_hidden_state.shape[1], inputs_dict["attention_mask"]
        )

        last_hidden_state_2 = decoder(
            input_ids=inputs_dict["decoder_input_ids"],
            attention_mask=inputs_dict["decoder_attention_mask"],
            encoder_hidden_states=encoder_last_hidden_state,
            encoder_attention_mask=encoder_attention_mask,
        )[0]

        self.parent.assertTrue((last_hidden_state_2 - last_hidden_state).abs().max().item() < 1e-3)


@require_torch
class Speech2TextModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (Speech2TextModel, Speech2TextForConditionalGeneration) if is_torch_available() else ()
    all_generative_model_classes = (Speech2TextForConditionalGeneration,) if is_torch_available() else ()
    pipeline_model_mapping = (
        {"automatic-speech-recognition": Speech2TextForConditionalGeneration, "feature-extraction": Speech2TextModel}
        if is_torch_available()
        else {}
    )
    is_encoder_decoder = True
    fx_compatible = True
    test_pruning = False
    test_missing_keys = False

    input_name = "input_features"

    def setUp(self):
        self.model_tester = Speech2TextModelTester(self)
        self.config_tester = ConfigTester(self, config_class=Speech2TextConfig)
        self.maxDiff = 3000

    def test_config(self):
        self.config_tester.run_common_tests()

    def test_save_load_strict(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
            model = model_class(config)

            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname)
                model2, info = model_class.from_pretrained(tmpdirname, output_loading_info=True)
            self.assertEqual(info["missing_keys"], [])

    def test_model_forward(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_forward(*config_and_inputs)

    def test_decoder_model_past_with_large_inputs(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)

    def test_encoder_decoder_model_standalone(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_encoder_decoder_model_standalone(*config_and_inputs)

    
    def test_inputs_embeds(self):
        pass

    
    def test_training(self):
        pass

    def test_training_gradient_checkpointing(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant(self):
        pass

    @unittest.skip(
        reason="This architecure seem to not compute gradients properly when using GC, check: https://github.com/huggingface/transformers/pull/27124"
    )
    def test_training_gradient_checkpointing_use_reentrant_false(self):
        pass

    @require_torch_fp16
    def test_generate_fp16(self):
        config, input_dict = self.model_tester.prepare_config_and_inputs()
        input_features = input_dict["input_features"]
        attention_mask = input_dict["attention_mask"]
        model = Speech2TextForConditionalGeneration(config).eval().to(torch_device)
        input_features = input_features.half()
        model.half()
        model.generate(input_features, attention_mask=attention_mask)
        model.generate(input_features, num_beams=4, do_sample=True, early_stopping=False, num_return_sequences=3)

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = [
                "input_features",
                "attention_mask",
                "decoder_input_ids",
                "decoder_attention_mask",
            ]
            expected_arg_names.extend(
                ["head_mask", "decoder_head_mask", "cross_attn_head_mask", "encoder_outputs"]
                if "head_mask" and "decoder_head_mask" and "cross_attn_head_mask" in arg_names
                else ["encoder_outputs"]
            )
            self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)

    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

            expected_num_layers = getattr(
                self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
            )
            self.assertEqual(len(hidden_states), expected_num_layers)

            if hasattr(self.model_tester, "encoder_seq_length"):
                seq_length = self.model_tester.encoder_seq_length
            else:
                seq_length = self.model_tester.seq_length

            subsampled_seq_length = model._get_feat_extract_output_lengths(seq_length)

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [subsampled_seq_length, self.model_tester.hidden_size],
            )

            if config.is_encoder_decoder:
                hidden_states = outputs.decoder_hidden_states

                self.assertIsInstance(hidden_states, (list, tuple))
                self.assertEqual(len(hidden_states), expected_num_layers)
                seq_len = getattr(self.model_tester, "seq_length", None)
                decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len)

                self.assertListEqual(
                    list(hidden_states[0].shape[-2:]),
                    [decoder_seq_length, self.model_tester.hidden_size],
                )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        seq_len = getattr(self.model_tester, "seq_length", None)
        decoder_seq_length = getattr(self.model_tester, "decoder_seq_length", seq_len)
        encoder_seq_length = getattr(self.model_tester, "encoder_seq_length", seq_len)
        decoder_key_length = getattr(self.model_tester, "decoder_key_length", decoder_seq_length)
        encoder_key_length = getattr(self.model_tester, "key_length", encoder_seq_length)

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            subsampled_encoder_seq_length = model._get_feat_extract_output_lengths(encoder_seq_length)
            subsampled_encoder_key_length = model._get_feat_extract_output_lengths(encoder_key_length)

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            self.assertListEqual(
                list(attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, subsampled_encoder_seq_length, subsampled_encoder_key_length],
            )
            out_len = len(outputs)

            correct_outlen = 5

            
            if "labels" in inputs_dict:
                correct_outlen += 1  
            if "past_key_values" in outputs:
                correct_outlen += 1  

            self.assertEqual(out_len, correct_outlen)

            
            decoder_attentions = outputs.decoder_attentions
            self.assertIsInstance(decoder_attentions, (list, tuple))
            self.assertEqual(len(decoder_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(decoder_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, decoder_seq_length, decoder_key_length],
            )

            
            cross_attentions = outputs.cross_attentions
            self.assertIsInstance(cross_attentions, (list, tuple))
            self.assertEqual(len(cross_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(cross_attentions[0].shape[-3:]),
                [
                    self.model_tester.num_attention_heads,
                    decoder_seq_length,
                    subsampled_encoder_key_length,
                ],
            )

            
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            added_hidden_states = 2
            self.assertEqual(out_len + added_hidden_states, len(outputs))

            self_attentions = outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

            self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads, subsampled_encoder_seq_length, subsampled_encoder_key_length],
            )

    def test_resize_tokens_embeddings(self):
        (
            original_config,
            inputs_dict,
        ) = self.model_tester.prepare_config_and_inputs_for_common()
        if not self.test_resize_embeddings:
            return

        for model_class in self.all_model_classes:
            config = copy.deepcopy(original_config)
            model = model_class(config)
            model.to(torch_device)

            if self.model_tester.is_training is False:
                model.eval()

            model_vocab_size = config.vocab_size
            
            model_embed = model.resize_token_embeddings(model_vocab_size)
            cloned_embeddings = model_embed.weight.clone()

            
            model_embed = model.resize_token_embeddings(model_vocab_size + 10)
            self.assertEqual(model.config.vocab_size, model_vocab_size + 10)
            
            self.assertEqual(model_embed.weight.shape[0], cloned_embeddings.shape[0] + 10)
            
            model(**self._prepare_for_class(inputs_dict, model_class))

            
            model_embed = model.resize_token_embeddings(model_vocab_size - 15)
            self.assertEqual(model.config.vocab_size, model_vocab_size - 15)
            
            self.assertEqual(model_embed.weight.shape[0], cloned_embeddings.shape[0] - 15)

            
            if "decoder_input_ids" in inputs_dict:
                inputs_dict["decoder_input_ids"].clamp_(max=model_vocab_size - 15 - 1)
            model(**self._prepare_for_class(inputs_dict, model_class))

            
            models_equal = True
            for p1, p2 in zip(cloned_embeddings, model_embed.weight):
                if p1.data.ne(p2.data).sum() > 0:
                    models_equal = False

            self.assertTrue(models_equal)

    def test_resize_embeddings_untied(self):
        (
            original_config,
            inputs_dict,
        ) = self.model_tester.prepare_config_and_inputs_for_common()
        if not self.test_resize_embeddings:
            return

        original_config.tie_word_embeddings = False

        
        if original_config.tie_word_embeddings:
            return

        for model_class in self.all_model_classes:
            config = copy.deepcopy(original_config)
            model = model_class(config).to(torch_device)

            
            if model.get_output_embeddings() is None:
                continue

            
            model_vocab_size = config.vocab_size
            model.resize_token_embeddings(model_vocab_size + 10)
            self.assertEqual(model.config.vocab_size, model_vocab_size + 10)
            output_embeds = model.get_output_embeddings()
            self.assertEqual(output_embeds.weight.shape[0], model_vocab_size + 10)
            
            if output_embeds.bias is not None:
                self.assertEqual(output_embeds.bias.shape[0], model_vocab_size + 10)
            
            model(**self._prepare_for_class(inputs_dict, model_class))

            
            model.resize_token_embeddings(model_vocab_size - 15)
            self.assertEqual(model.config.vocab_size, model_vocab_size - 15)
            
            output_embeds = model.get_output_embeddings()
            self.assertEqual(output_embeds.weight.shape[0], model_vocab_size - 15)
            
            if output_embeds.bias is not None:
                self.assertEqual(output_embeds.bias.shape[0], model_vocab_size - 15)
            
            if "decoder_input_ids" in inputs_dict:
                inputs_dict["decoder_input_ids"].clamp_(max=model_vocab_size - 15 - 1)
            
            model(**self._prepare_for_class(inputs_dict, model_class))

    def test_generate_without_input_ids(self):
        pass

    @staticmethod
    def _get_encoder_outputs(
        model, input_ids, attention_mask, output_attentions=None, output_hidden_states=None, num_interleave=1
    ):
        encoder = model.get_encoder()
        encoder_outputs = encoder(
            input_ids,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
        )
        encoder_outputs["last_hidden_state"] = encoder_outputs.last_hidden_state.repeat_interleave(
            num_interleave, dim=0
        )
        input_ids = input_ids[:, :, 0]
        input_ids = torch.zeros_like(input_ids[:, :1], dtype=torch.long) + model._get_decoder_start_token_id()
        attention_mask = None
        return encoder_outputs, input_ids, attention_mask

    def _check_outputs(self, output, input_ids, config, use_cache=False, num_return_sequences=1):
        batch_size, seq_length = input_ids.shape[:2]
        subsampled_seq_length = self.model_tester.get_subsampled_output_lengths(seq_length)
        num_sequences_in_output = batch_size * num_return_sequences
        gen_len = (
            output.sequences.shape[-1] - 1 if config.is_encoder_decoder else output.sequences.shape[-1] - seq_length
        )

        
        self._check_scores(num_sequences_in_output, output.scores, length=gen_len, config=config)

        
        
        self._check_encoder_attention_for_generate(
            output.encoder_attentions, batch_size, config, subsampled_seq_length
        )
        
        self._check_attentions_for_generate(
            num_sequences_in_output,
            output.decoder_attentions,
            min_length=1,
            max_length=output.sequences.shape[-1],
            config=config,
            use_cache=use_cache,
        )

        
        
        self._check_encoder_hidden_states_for_generate(
            output.encoder_hidden_states, batch_size, config, subsampled_seq_length
        )

        
        self._check_hidden_states_for_generate(
            num_sequences_in_output,
            output.decoder_hidden_states,
            min_length=1,
            max_length=output.sequences.shape[-1],
            config=config,
            use_cache=use_cache,
        )

    def _create_and_check_torchscript(self, config, inputs_dict):
        if not self.test_torchscript:
            return

        configs_no_init = _config_zero_init(config)  
        configs_no_init.torchscript = True
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            model.to(torch_device)
            model.eval()
            inputs = self._prepare_for_class(inputs_dict, model_class)

            try:
                model.config.use_cache = False  
                input_features = inputs["input_features"]
                attention_mask = inputs["attention_mask"]
                decoder_input_ids = inputs["decoder_input_ids"]
                decoder_attention_mask = inputs["decoder_attention_mask"]
                traced_model = torch.jit.trace(
                    model, (input_features, attention_mask, decoder_input_ids, decoder_attention_mask)
                )
            except RuntimeError:
                self.fail("Couldn't trace module.")

            with tempfile.TemporaryDirectory() as tmp_dir_name:
                pt_file_name = os.path.join(tmp_dir_name, "traced_model.pt")

                try:
                    torch.jit.save(traced_model, pt_file_name)
                except Exception:
                    self.fail("Couldn't save module.")

                try:
                    loaded_model = torch.jit.load(pt_file_name)
                except Exception:
                    self.fail("Couldn't load module.")

            model.to(torch_device)
            model.eval()

            loaded_model.to(torch_device)
            loaded_model.eval()

            model_state_dict = model.state_dict()
            loaded_model_state_dict = loaded_model.state_dict()

            non_persistent_buffers = {}
            for key in loaded_model_state_dict.keys():
                if key not in model_state_dict.keys():
                    non_persistent_buffers[key] = loaded_model_state_dict[key]

            loaded_model_state_dict = {
                key: value for key, value in loaded_model_state_dict.items() if key not in non_persistent_buffers
            }

            self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys()))

            model_buffers = list(model.buffers())
            for non_persistent_buffer in non_persistent_buffers.values():
                found_buffer = False
                for i, model_buffer in enumerate(model_buffers):
                    if torch.equal(non_persistent_buffer, model_buffer):
                        found_buffer = True
                        break

                self.assertTrue(found_buffer)
                model_buffers.pop(i)

            models_equal = True
            for layer_name, p1 in model_state_dict.items():
                p2 = loaded_model_state_dict[layer_name]
                if p1.data.ne(p2.data).sum() > 0:
                    models_equal = False

            self.assertTrue(models_equal)

    def test_pt_tf_model_equivalence(self, allow_missing_keys=True):
        
        super().test_pt_tf_model_equivalence(allow_missing_keys=allow_missing_keys)

    @unittest.skip("Test failing,  @RocketNight is looking into it")
    def test_tf_from_pt_safetensors(self):
        pass


@require_torch
@require_torchaudio
@require_sentencepiece
@require_tokenizers
@slow
class Speech2TextModelIntegrationTests(unittest.TestCase):
    @cached_property
    def default_processor(self):
        return Speech2TextProcessor.from_pretrained("facebook/s2t-small-librispeech-asr")

    def _load_datasamples(self, num_samples):
        from datasets import load_dataset

        ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
        
        speech_samples = ds.sort("id").select(range(num_samples))[:num_samples]["audio"]

        return [x["array"] for x in speech_samples]

    def test_generation_librispeech(self):
        model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-small-librispeech-asr")
        model.to(torch_device)
        processor = self.default_processor

        input_speech = self._load_datasamples(1)

        input_features = processor(input_speech, return_tensors="pt").input_features.to(torch_device)

        generated_ids = model.generate(input_features)
        generated_transcript = processor.batch_decode(generated_ids, skip_special_tokens=True)

        EXPECTED_TRANSCRIPTIONS = [
            "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel"
        ]
        self.assertListEqual(generated_transcript, EXPECTED_TRANSCRIPTIONS)

    def test_generation_librispeech_batched(self):
        model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-small-librispeech-asr")
        model.to(torch_device)
        processor = self.default_processor

        input_speech = self._load_datasamples(4)

        inputs = processor(input_speech, return_tensors="pt", padding=True)

        input_features = inputs.input_features.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        generated_ids = model.generate(input_features, attention_mask=attention_mask)
        generated_transcripts = processor.batch_decode(generated_ids, skip_special_tokens=True)

        EXPECTED_TRANSCRIPTIONS = [
            "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel",
            "nor is mister cultar's manner less interesting than his matter",
            "he tells us that at this festive season of the year with christmas and roast beef looming before us"
            " similes drawn from eating and its results occur most readily to the mind",
            "he has grave doubts whether sir frederick leyton's work is really greek after all and can discover in it"
            " but little of rocky ithaca",
        ]

        self.assertListEqual(generated_transcripts, EXPECTED_TRANSCRIPTIONS)