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 hubert model import math import os import pickle import tempfile import unittest import pytest from transformers import hubertconfig istorchavailable from transformers testingutils import requiresoundfile requiretorch slow torchdevice from transformers utils import istorchfxavailable from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin configzeroinit floatstensor idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import hubertforctc hubertforsequenceclassification hubertmodel wav2vec2featureextractor wav2vec2processor from transformers models hubert modelinghubert import computemaskindices if istorchfxavailable from transformers utils fx import symbolictrace class hubertmodeltester def init self parent batchsize13 seqlength1024 speech is longer istrainingfalse hiddensize16 featextractnormgroup featextractdropout0 0 featextractactivationgelu convdim32 32 32 convstride4 4 4 convkernel8 8 8 convbiasfalse numconvposembeddings16 numconvposembeddinggroups2 numhiddenlayers2 numattentionheads2 hiddendropoutprob0 1 this is most likely not correctly set yet intermediatesize20 layernormeps1e5 hiddenactgelu initializerrange0 02 vocabsize32 dostablelayernormfalse scopenone self parent parent self batchsize batchsize self seqlength seqlength self istraining istraining self hiddensize hiddensize self featextractnorm featextractnorm self featextractdropout featextractdropout self featextractactivation featextractactivation self convdim convdim self convstride convstride self convkernel convkernel self convbias convbias self numconvposembeddings numconvposembeddings self numconvposembeddinggroups numconvposembeddinggroups self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self hiddendropoutprob hiddendropoutprob self intermediatesize intermediatesize self layernormeps layernormeps self hiddenact hiddenact self initializerrange initializerrange self vocabsize vocabsize self dostablelayernorm dostablelayernorm self scope scope outputseqlength self seqlength for kernel stride in zipself convkernel self convstride outputseqlength outputseqlength kernel 1 stride self outputseqlength intmath ceiloutputseqlength self encoderseqlength self outputseqlength def prepareconfigandinputsself inputvalues floatstensorself batchsize self seqlength scale1 0 attentionmask randomattentionmaskself batchsize self seqlength config self getconfig return config inputvalues attentionmask def getconfigself return hubertconfig hiddensizeself hiddensize featextractnormself featextractnorm featextractdropoutself featextractdropout featextractactivationself featextractactivation convdimself convdim convstrideself convstride convkernelself convkernel convbiasself convbias numconvposembeddingsself numconvposembeddings numconvposembeddinggroupsself numconvposembeddinggroups numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads hiddendropoutprobself hiddendropoutprob intermediatesizeself intermediatesize layernormepsself layernormeps hiddenactself hiddenact initializerrangeself initializerrange vocabsizeself vocabsize dostablelayernormself dostablelayernorm def createandcheckmodelself config inputvalues attentionmask model hubertmodelconfigconfig model totorchdevice model eval result modelinputvalues attentionmaskattentionmask self parent assertequal result lasthiddenstate shape self batchsize self outputseqlength self hiddensize def createandcheckbatchinferenceself config inputvalues args test does not pass for models making use of groupnorm check https github compytorchfairseqissues3227 model hubertmodelconfigconfig model totorchdevice model eval inputvalues inputvalues 3 attentionmask torch onesinputvalues shape devicetorchdevice dtypetorch bool inputlengths inputvalues shape1 i for i in 4 2 1 pad input for i in rangeleninputlengths inputvaluesi inputlengthsi 0 0 attentionmaski inputlengthsi 0 0 batchoutputs modelinputvalues attentionmaskattentionmask lasthiddenstate for i in rangeinputvalues shape0 inputslice inputvaluesi i 1 inputlengthsi output modelinputslice lasthiddenstate batchoutput batchoutputsi i 1 output shape1 self parent asserttruetorch allcloseoutput batchoutput atol1e3 def checkctclossself config inputvalues args model hubertforctcconfigconfig model totorchdevice make sure that dropout is disabled model eval inputvalues inputvalues 3 attentionmask torch onesinputvalues shape devicetorchdevice dtypetorch long inputlengths inputvalues shape1 i for i in 4 2 1 maxlengthlabels model getfeatextractoutputlengthstorch tensorinputlengths labels idstensorinputvalues shape0 minmaxlengthlabels 1 model config vocabsize pad input for i in rangeleninputlengths inputvaluesi inputlengthsi 0 0 attentionmaski inputlengthsi 0 model config ctclossreduction sum sumloss modelinputvalues attentionmaskattentionmask labelslabels loss item model config ctclossreduction mean meanloss modelinputvalues attentionmaskattentionmask labelslabels loss item self parent asserttrueisinstancesumloss float self parent asserttrueisinstancemeanloss float def checkseqclassifierlossself config inputvalues args model hubertforsequenceclassificationconfigconfig model totorchdevice make sure that dropout is disabled model eval inputvalues inputvalues 3 attentionmask torch onesinputvalues shape devicetorchdevice dtypetorch long inputlengths inputvalues shape1 i for i in 4 2 1 labels idstensorinputvalues shape0 1 lenmodel config id2label pad input for i in rangeleninputlengths inputvaluesi inputlengthsi 0 0 attentionmaski inputlengthsi 0 maskedloss modelinputvalues attentionmaskattentionmask labelslabels loss item unmaskedloss modelinputvalues labelslabels loss item self parent asserttrueisinstancemaskedloss float self parent asserttrueisinstanceunmaskedloss float self parent asserttruemaskedloss unmaskedloss def checkctctrainingself config inputvalues args config ctczeroinfinity true model hubertforctcconfigconfig model totorchdevice model train freeze feature encoder model freezefeatureencoder inputvalues inputvalues 3 inputlengths inputvalues shape1 i for i in 4 2 1 maxlengthlabels model getfeatextractoutputlengthstorch tensorinputlengths labels idstensorinputvalues shape0 maxmaxlengthlabels 2 model config vocabsize pad input for i in rangeleninputlengths inputvaluesi inputlengthsi 0 0 if maxlengthlabelsi labels shape1 it s important that we make sure that target lengths are at least one shorter than logit lengths to prevent inf labelsi maxlengthlabelsi 1 100 loss modelinputvalues labelslabels loss self parent assertfalsetorch isinfloss item loss backward def checkseqclassifiertrainingself config inputvalues args config ctczeroinfinity true model hubertforsequenceclassificationconfigconfig model totorchdevice model train freeze everything but the classification head model freezebasemodel inputvalues inputvalues 3 inputlengths inputvalues shape1 i for i in 4 2 1 labels idstensorinputvalues shape0 1 lenmodel config id2label pad input for i in rangeleninputlengths inputvaluesi inputlengthsi 0 0 loss modelinputvalues labelslabels loss self parent assertfalsetorch isinfloss item loss backward def checklabelsoutofvocabself config inputvalues args model hubertforctcconfig model totorchdevice model train inputvalues inputvalues 3 inputlengths inputvalues shape1 i for i in 4 2 1 maxlengthlabels model getfeatextractoutputlengthstorch tensorinputlengths labels idstensorinputvalues shape0 maxmaxlengthlabels 2 model config vocabsize 100 with pytest raisesvalueerror modelinputvalues labelslabels def prepareconfigandinputsforcommonself config inputvalues attentionmask self prepareconfigandinputs inputsdict inputvalues inputvalues attentionmask attentionmask return config inputsdict requiretorch class hubertmodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses hubertforctc hubertforsequenceclassification hubertmodel if istorchavailable else pipelinemodelmapping audioclassification hubertforsequenceclassification automaticspeechrecognition hubertforctc featureextraction hubertmodel if istorchavailable else fxcompatible true testpruning false testheadmasking false def setupself self modeltester hubertmodeltesterself self configtester configtesterself configclasshubertconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testctclossinferenceself configandinputs self modeltester prepareconfigandinputs self modeltester checkctclossconfigandinputs def testseqclassifierlossinferenceself configandinputs self modeltester prepareconfigandinputs self modeltester checkseqclassifierlossconfigandinputs def testctctrainself configandinputs self modeltester prepareconfigandinputs self modeltester checkctctrainingconfigandinputs def testseqclassifiertrainself configandinputs self modeltester prepareconfigandinputs self modeltester checkseqclassifiertrainingconfigandinputs def testlabelsoutofvocabself configandinputs self modeltester prepareconfigandinputs self modeltester checklabelsoutofvocabconfigandinputs hubert has no inputsembeds def testinputsembedsself pass inputids is renamed to inputvalues def testforwardsignatureself pass hubert cannot resize token embeddings since it has no tokens embeddings def testresizetokensembeddingsself pass hubert has no inputsembeds and thus the getinputembeddings fn is not implemented def testmodelcommonattributesself pass def testretaingradhiddenstatesattentionsself config inputsdict self modeltester prepareconfigandinputsforcommon config outputhiddenstates true config outputattentions true no need to test all models as different heads yield the same functionality modelclass self allmodelclasses0 model modelclassconfig model totorchdevice set layer drop to 0 model config layerdrop 0 0 inputvalues inputsdictinputvalues inputlengths torch tensor inputvalues shape1 for in rangeinputvalues shape0 dtypetorch long devicetorchdevice outputlengths model getfeatextractoutputlengthsinputlengths labels idstensorinputvalues shape0 outputlengths0 2 self modeltester vocabsize inputsdictattentionmask torch oneslikeinputsdictattentionmask inputsdictlabels labels outputs modelinputsdict output outputs0 encoderdecoderonly models hiddenstates outputs hiddenstates0 attentions outputs attentions0 hiddenstates retaingrad attentions retaingrad output flatten0 backwardretaingraphtrue self assertisnotnonehiddenstates grad self assertisnotnoneattentions grad def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters uniforminitparms conv weight conv parametrizations weight maskedspecembed quantizer weightproj weight if param requiresgrad if anyx in name for x in uniforminitparms self asserttrue 1 0 param data mean 1e9 round 1e9 item 1 0 msgfparameter name of model modelclass seems not properly initialized else self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized hubert cannot be torchscripted because of torch nn utils weightnorm def createandchecktorchfxtracingself config inputsdict outputlossfalse todo fix it self skiptesttorch 2 1 breaks torch fx tests for wav2vec2hubert if not istorchfxavailable or not self fxcompatible return configsnoinit configzeroinitconfig to be sure we have no nan configsnoinit returndict false for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit model totorchdevice model eval inputs self prepareforclassinputsdict modelclass returnlabelsoutputloss try if model config isencoderdecoder model config usecache false fstm still requires this hack fstm should probably be refactored similar to bart afterward labels inputs getlabels none inputnames attentionmask decoderattentionmask decoderinputids inputfeatures inputids inputvalues if labels is not none inputnames appendlabels filteredinputs k v for k v in inputs items if k in inputnames inputnames listfilteredinputs keys modeloutput modelfilteredinputs tracedmodel symbolictracemodel inputnames tracedoutput tracedmodelfilteredinputs else inputnames attentionmask bbox inputfeatures inputids inputvalues pixelvalues tokentypeids visualfeats visualpos labels inputs getlabels none startpositions inputs getstartpositions none endpositions inputs getendpositions none if labels is not none inputnames appendlabels if startpositions is not none inputnames appendstartpositions if endpositions is not none inputnames appendendpositions filteredinputs k v for k v in inputs items if k in inputnames inputnames listfilteredinputs keys modeloutput modelfilteredinputs tracedmodel symbolictracemodel inputnames tracedoutput tracedmodelfilteredinputs except exception as e self failfcouldn t trace module e def flattenoutputoutput flatten for x in output if isinstancex tuple list flatten flattenoutputx elif not isinstancex torch tensor continue else flatten appendx return flatten modeloutput flattenoutputmodeloutput tracedoutput flattenoutputtracedoutput numoutputs lenmodeloutput for i in rangenumoutputs self asserttrue torch allclosemodeloutputi tracedoutputi ftraced ith output doesn t match model ith output for modelclass test that the model can be serialized and restored properly with tempfile temporarydirectory as tmpdirname pklfilename os path jointmpdirname model pkl try with openpklfilename wb as f pickle dumptracedmodel f with openpklfilename rb as f loaded pickle loadf except exception as e self failfcouldn t serialize deserialize the traced model e loadedoutput loadedfilteredinputs loadedoutput flattenoutputloadedoutput for i in rangenumoutputs self asserttrue torch allclosemodeloutputi loadedoutputi fserialized model ith output doesn t match model ith output for modelclass overwrite from testmodelingcommon def mockinitweightsself module if hasattrmodule weight and module weight is not none module weight data fill3 if hasattrmodule weightg and module weightg is not none module weightg data fill3 if hasattrmodule weightv and module weightv is not none module weightv data fill3 if hasattrmodule bias and module bias is not none module bias data fill3 if hasattrmodule maskedspecembed and module maskedspecembed is not none module maskedspecembed data fill3 unittest skipreasonfeed forward chunking is not implemented def testfeedforwardchunkingself pass slow def testmodelfrompretrainedself model hubertmodel frompretrainedfacebookhubertbasels960 self assertisnotnonemodel requiretorch class hubertrobustmodeltestmodeltestermixin unittest testcase allmodelclasses hubertforctc hubertforsequenceclassification hubertmodel if istorchavailable else testpruning false testheadmasking false def setupself self modeltester hubertmodeltester self convstride3 3 3 featextractnormlayer dostablelayernormtrue self configtester configtesterself configclasshubertconfig hiddensize37 def testconfigself self configtester runcommontests def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testbatchedinferenceself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckbatchinferenceconfigandinputs def testctclossinferenceself configandinputs self modeltester prepareconfigandinputs self modeltester checkctclossconfigandinputs def testseqclassifierlossinferenceself configandinputs self modeltester prepareconfigandinputs self modeltester checkseqclassifierlossconfigandinputs def testctctrainself configandinputs self modeltester prepareconfigandinputs self modeltester checkctctrainingconfigandinputs def testseqclassifiertrainself configandinputs self modeltester prepareconfigandinputs self modeltester checkseqclassifiertrainingconfigandinputs def testlabelsoutofvocabself configandinputs self modeltester prepareconfigandinputs self modeltester checklabelsoutofvocabconfigandinputs hubert has no inputsembeds def testinputsembedsself pass inputids is renamed to inputvalues def testforwardsignatureself pass hubert cannot resize token embeddings since it has no tokens embeddings def testresizetokensembeddingsself pass hubert has no inputsembeds and thus the getinputembeddings fn is not implemented def testmodelcommonattributesself pass def testretaingradhiddenstatesattentionsself config inputsdict self modeltester prepareconfigandinputsforcommon config outputhiddenstates true config outputattentions true no need to test all models as different heads yield the same functionality modelclass self allmodelclasses0 model modelclassconfig model totorchdevice set layer drop to 0 model config layerdrop 0 0 inputvalues inputsdictinputvalues inputlengths torch tensor inputvalues shape1 for in rangeinputvalues shape0 dtypetorch long devicetorchdevice outputlengths model getfeatextractoutputlengthsinputlengths labels idstensorinputvalues shape0 outputlengths0 2 self modeltester vocabsize inputsdictattentionmask torch oneslikeinputsdictattentionmask inputsdictlabels labels outputs modelinputsdict output outputs0 encoderdecoderonly models hiddenstates outputs hiddenstates0 attentions outputs attentions0 hiddenstates retaingrad attentions retaingrad output flatten0 backwardretaingraphtrue self assertisnotnonehiddenstates grad self assertisnotnoneattentions grad def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters uniforminitparms conv weight conv parametrizations weight maskedspecembed quantizer weightproj weight if param requiresgrad if anyx in name for x in uniforminitparms self asserttrue 1 0 param data mean 1e9 round 1e9 item 1 0 msgfparameter name of model modelclass seems not properly initialized else self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized overwrite from testmodelingcommon def mockinitweightsself module if hasattrmodule weight and module weight is not none module weight data fill3 if hasattrmodule weightg and module weightg is not none module weightg data fill3 if hasattrmodule weightv and module weightv is not none module weightv data fill3 if hasattrmodule bias and module bias is not none module bias data fill3 if hasattrmodule maskedspecembed and module maskedspecembed is not none module maskedspecembed data fill3 unittest skipreasonfeed forward chunking is not implemented def testfeedforwardchunkingself pass slow def testmodelfrompretrainedself model hubertmodel frompretrainedfacebookhubertlargels960ft self assertisnotnonemodel requiretorch class hubertutilstestunittest testcase def testcomputemaskindicesself batchsize 4 sequencelength 60 maskprob 0 5 masklength 1 mask computemaskindicesbatchsize sequencelength maskprob masklength mask torch fromnumpymask totorchdevice self assertlistequalmask sumaxis1 tolist maskprob sequencelength for in rangebatchsize def testcomputemaskindicesoverlapself batchsize 4 sequencelength 80 maskprob 0 5 masklength 4 mask computemaskindicesbatchsize sequencelength maskprob masklength mask torch fromnumpymask totorchdevice because of overlap mask don t have to add up exactly to maskprob sequencelength but have to be smaller or equal for batchsum in mask sumaxis1 self asserttrueintbatchsum maskprob sequencelength requiretorch requiresoundfile slow class hubertmodelintegrationtestunittest testcase def loaddatasamplesself numsamples from datasets import loaddataset ds loaddatasethfinternaltestinglibrispeechasrdummy clean splitvalidation automatic decoding with librispeech speechsamples ds sortid filter lambda x xid in f1272141231000i for i in rangenumsamples numsamplesaudio return xarray for x in speechsamples def loadsuperbself task numsamples from datasets import loaddataset ds loaddatasetantonlsuperbdummy task splittest return ds numsamples def testinferencectcbatchedself model hubertforctc frompretrainedfacebookhubertlargels960ft torchdtypetorch float16 to torchdevice processor wav2vec2processor frompretrainedfacebookhubertlargels960ft dolowercasetrue inputspeech self loaddatasamples2 inputs processorinputspeech returntensorspt paddingtrue inputvalues inputs inputvalues half totorchdevice attentionmask inputs attentionmask totorchdevice with torch nograd logits modelinputvalues attentionmaskattentionmask logits predictedids torch argmaxlogits dim1 predictedtrans processor batchdecodepredictedids expectedtranscriptions a man said to the universe sir i exist sweat covered brion s body trickling into the tight loin cloth that was the only garment he wore self assertlistequalpredictedtrans expectedtranscriptions def testinferencekeywordspottingself model hubertforsequenceclassification frompretrained superbhubertbasesuperbks torchdtypetorch float16 totorchdevice processor wav2vec2featureextractor frompretrainedsuperbhubertbasesuperbks inputdata self loadsuperbks 4 inputs processorinputdataspeech returntensorspt paddingtrue inputvalues inputs inputvalues half totorchdevice attentionmask inputs attentionmask totorchdevice with torch nograd outputs modelinputvalues attentionmaskattentionmask predictedlogits predictedids torch maxoutputs logits dim1 expectedlabels 2 6 10 9 s3prl logits for the same batch expectedlogits torch tensor7 6692 17 7795 11 1562 11 8232 dtypetorch float16 devicetorchdevice self assertlistequalpredictedids tolist expectedlabels self asserttruetorch allclosepredictedlogits expectedlogits atol3e2 def testinferenceintentclassificationself model hubertforsequenceclassification frompretrained superbhubertbasesuperbic torchdtypetorch float16 totorchdevice processor wav2vec2featureextractor frompretrainedsuperbhubertbasesuperbic inputdata self loadsuperbic 4 inputs processorinputdataspeech returntensorspt paddingtrue inputvalues inputs inputvalues half totorchdevice attentionmask inputs attentionmask totorchdevice with torch nograd outputs modelinputvalues attentionmaskattentionmask predictedlogitsaction predictedidsaction torch maxoutputs logits 6 dim1 predictedlogitsobject predictedidsobject torch maxoutputs logits 6 20 dim1 predictedlogitslocation predictedidslocation torch maxoutputs logits 20 24 dim1 expectedlabelsaction 1 0 4 3 expectedlogitsaction torch tensor 5 9052 12 5865 4 4840 10 0240 dtypetorch float16 devicetorchdevice expectedlabelsobject 1 10 3 4 expectedlogitsobject torch tensor 5 5316 11 7946 8 1672 23 2415 dtypetorch float16 devicetorchdevice expectedlabelslocation 0 0 0 1 expectedlogitslocation torch tensor 5 2053 8 9577 10 0447 8 1481 dtypetorch float16 devicetorchdevice self assertlistequalpredictedidsaction tolist expectedlabelsaction self assertlistequalpredictedidsobject tolist expectedlabelsobject self assertlistequalpredictedidslocation tolist expectedlabelslocation todo lower the tolerance after merging the padding fix https github compytorchfairseqpull3572 self asserttruetorch allclosepredictedlogitsaction expectedlogitsaction atol3e1 self asserttruetorch allclosepredictedlogitsobject expectedlogitsobject atol3e1 self asserttruetorch allclosepredictedlogitslocation expectedlogitslocation atol3e1 def testinferencespeakeridentificationself model hubertforsequenceclassification frompretrained superbhubertbasesuperbsid torchdtypetorch float16 totorchdevice processor wav2vec2featureextractor frompretrainedsuperbhubertbasesuperbsid inputdata self loadsuperbsi 4 outputlogits with torch nograd for example in inputdataspeech input processorexample returntensorspt paddingtrue output modelinput inputvalues half totorchdevice attentionmasknone outputlogits appendoutput logits0 outputlogits torch stackoutputlogits predictedlogits predictedids torch maxoutputlogits dim1 expectedlabels 5 1 1 3 s3prl logits for the same batch expectedlogits torch tensor 78231 5547 123166 6094 122785 4141 84851 2969 dtypetorch float16 devicetorchdevice self assertlistequalpredictedids tolist expectedlabels todo lower the tolerance after merging the padding fix https github compytorchfairseqpull3572 self asserttruetorch allclosepredictedlogits expectedlogits atol10 def testinferenceemotionrecognitionself model hubertforsequenceclassification frompretrained superbhubertbasesuperber torchdtypetorch float16 totorchdevice processor wav2vec2featureextractor frompretrainedsuperbhubertbasesuperber inputdata self loadsuperber 4 inputs processorinputdataspeech returntensorspt paddingtrue inputvalues inputs inputvalues half totorchdevice attentionmask inputs attentionmask totorchdevice with torch nograd outputs modelinputvalues attentionmaskattentionmask predictedlogits predictedids torch maxoutputs logits dim1 expectedlabels 1 1 2 2 s3prl logits for the same batch expectedlogits torch tensor2 8384 2 3389 3 8564 4 5558 dtypetorch float16 devicetorchdevice self assertlistequalpredictedids tolist expectedlabels todo lower the tolerance after merging the padding fix https github compytorchfairseqpull3572 self asserttruetorch allclosepredictedlogits expectedlogits atol1e1 def testinferencedistilhubertself model hubertmodel frompretrainedntuspmldistilhubert totorchdevice processor wav2vec2featureextractor frompretrainedntuspmldistilhubert todo can t test on batched inputs due to incompatible padding https github compytorchfairseqpull3572 inputspeech self loaddatasamples1 inputs processorinputspeech returntensorspt paddingtrue inputvalues inputs inputvalues totorchdevice with torch nograd outputs modelinputvalues lasthiddenstate expected outputs taken from the original sew implementation expectedoutputsfirst torch tensor 0 3505 0 1167 0 0608 0 1294 0 3085 0 0481 0 1106 0 0955 0 3107 0 0391 0 0739 0 1360 0 2385 0 1795 0 0928 0 2389 devicetorchdevice expectedoutputslast torch tensor 0 0732 0 0255 0 0529 0 1372 0 0812 0 1259 0 0564 0 0438 0 0054 0 0758 0 0002 0 1617 0 0133 0 0320 0 0687 0 0062 devicetorchdevice expectedoutputsum 3776 0730 self asserttruetorch allcloseoutputs 4 4 expectedoutputsfirst atol5e3 self asserttruetorch allcloseoutputs 4 4 expectedoutputslast atol5e3 self asserttrueabsoutputs sum expectedoutputsum 0 1 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 hubert model speech is longer this is most likely not correctly set yet test does not pass for models making use of group_norm check https github com pytorch fairseq issues 3227 pad input make sure that dropout is disabled pad input make sure that dropout is disabled pad input freeze feature encoder pad input it s important that we make sure that target lengths are at least one shorter than logit lengths to prevent inf freeze everything but the classification head pad input hubert has no inputs_embeds input_ids is renamed to input_values hubert cannot resize token embeddings since it has no tokens embeddings hubert has no inputs_embeds and thus the get_input_embeddings fn is not implemented no need to test all models as different heads yield the same functionality set layer drop to 0 encoder decoder only models hubert cannot be torchscripted because of torch nn utils weight_norm todo fix it to be sure we have no nan fstm still requires this hack fstm should probably be refactored similar to bart afterward test that the model can be serialized and restored properly overwrite from test_modeling_common hubert has no inputs_embeds input_ids is renamed to input_values hubert cannot resize token embeddings since it has no tokens embeddings hubert has no inputs_embeds and thus the get_input_embeddings fn is not implemented no need to test all models as different heads yield the same functionality set layer drop to 0 encoder decoder only models overwrite from test_modeling_common because of overlap mask don t have to add up exactly to mask_prob sequence_length but have to be smaller or equal automatic decoding with librispeech s3prl logits for the same batch todo lower the tolerance after merging the padding fix https github com pytorch fairseq pull 3572 s3prl logits for the same batch todo lower the tolerance after merging the padding fix https github com pytorch fairseq pull 3572 s3prl logits for the same batch todo lower the tolerance after merging the padding fix https github com pytorch fairseq pull 3572 todo can t test on batched inputs due to incompatible padding https github com pytorch fairseq pull 3572 expected outputs taken from the original sew implementation
import math
import os
import pickle
import tempfile
import unittest

import pytest

from transformers import HubertConfig, is_torch_available
from transformers.testing_utils import require_soundfile, require_torch, slow, torch_device
from transformers.utils import is_torch_fx_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import (
    ModelTesterMixin,
    _config_zero_init,
    floats_tensor,
    ids_tensor,
    random_attention_mask,
)
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        HubertForCTC,
        HubertForSequenceClassification,
        HubertModel,
        Wav2Vec2FeatureExtractor,
        Wav2Vec2Processor,
    )
    from transformers.models.hubert.modeling_hubert import _compute_mask_indices

if is_torch_fx_available():
    from transformers.utils.fx import symbolic_trace


class HubertModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        seq_length=1024,  
        is_training=False,
        hidden_size=16,
        feat_extract_norm="group",
        feat_extract_dropout=0.0,
        feat_extract_activation="gelu",
        conv_dim=(32, 32, 32),
        conv_stride=(4, 4, 4),
        conv_kernel=(8, 8, 8),
        conv_bias=False,
        num_conv_pos_embeddings=16,
        num_conv_pos_embedding_groups=2,
        num_hidden_layers=2,
        num_attention_heads=2,
        hidden_dropout_prob=0.1,  
        intermediate_size=20,
        layer_norm_eps=1e-5,
        hidden_act="gelu",
        initializer_range=0.02,
        vocab_size=32,
        do_stable_layer_norm=False,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.hidden_size = hidden_size
        self.feat_extract_norm = feat_extract_norm
        self.feat_extract_dropout = feat_extract_dropout
        self.feat_extract_activation = feat_extract_activation
        self.conv_dim = conv_dim
        self.conv_stride = conv_stride
        self.conv_kernel = conv_kernel
        self.conv_bias = conv_bias
        self.num_conv_pos_embeddings = num_conv_pos_embeddings
        self.num_conv_pos_embedding_groups = num_conv_pos_embedding_groups
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_dropout_prob = hidden_dropout_prob
        self.intermediate_size = intermediate_size
        self.layer_norm_eps = layer_norm_eps
        self.hidden_act = hidden_act
        self.initializer_range = initializer_range
        self.vocab_size = vocab_size
        self.do_stable_layer_norm = do_stable_layer_norm
        self.scope = scope

        output_seq_length = self.seq_length
        for kernel, stride in zip(self.conv_kernel, self.conv_stride):
            output_seq_length = (output_seq_length - (kernel - 1)) / stride
        self.output_seq_length = int(math.ceil(output_seq_length))
        self.encoder_seq_length = self.output_seq_length

    def prepare_config_and_inputs(self):
        input_values = floats_tensor([self.batch_size, self.seq_length], scale=1.0)
        attention_mask = random_attention_mask([self.batch_size, self.seq_length])

        config = self.get_config()

        return config, input_values, attention_mask

    def get_config(self):
        return HubertConfig(
            hidden_size=self.hidden_size,
            feat_extract_norm=self.feat_extract_norm,
            feat_extract_dropout=self.feat_extract_dropout,
            feat_extract_activation=self.feat_extract_activation,
            conv_dim=self.conv_dim,
            conv_stride=self.conv_stride,
            conv_kernel=self.conv_kernel,
            conv_bias=self.conv_bias,
            num_conv_pos_embeddings=self.num_conv_pos_embeddings,
            num_conv_pos_embedding_groups=self.num_conv_pos_embedding_groups,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            hidden_dropout_prob=self.hidden_dropout_prob,
            intermediate_size=self.intermediate_size,
            layer_norm_eps=self.layer_norm_eps,
            hidden_act=self.hidden_act,
            initializer_range=self.initializer_range,
            vocab_size=self.vocab_size,
            do_stable_layer_norm=self.do_stable_layer_norm,
        )

    def create_and_check_model(self, config, input_values, attention_mask):
        model = HubertModel(config=config)
        model.to(torch_device)
        model.eval()
        result = model(input_values, attention_mask=attention_mask)
        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, self.output_seq_length, self.hidden_size)
        )

    def create_and_check_batch_inference(self, config, input_values, *args):
        
        
        model = HubertModel(config=config)
        model.to(torch_device)
        model.eval()

        input_values = input_values[:3]
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.bool)

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]

        
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
            attention_mask[i, input_lengths[i] :] = 0.0

        batch_outputs = model(input_values, attention_mask=attention_mask).last_hidden_state

        for i in range(input_values.shape[0]):
            input_slice = input_values[i : i + 1, : input_lengths[i]]
            output = model(input_slice).last_hidden_state

            batch_output = batch_outputs[i : i + 1, : output.shape[1]]
            self.parent.assertTrue(torch.allclose(output, batch_output, atol=1e-3))

    def check_ctc_loss(self, config, input_values, *args):
        model = HubertForCTC(config=config)
        model.to(torch_device)

        
        model.eval()

        input_values = input_values[:3]
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.long)

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], min(max_length_labels) - 1), model.config.vocab_size)

        
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
            attention_mask[i, input_lengths[i] :] = 0

        model.config.ctc_loss_reduction = "sum"
        sum_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()

        model.config.ctc_loss_reduction = "mean"
        mean_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()

        self.parent.assertTrue(isinstance(sum_loss, float))
        self.parent.assertTrue(isinstance(mean_loss, float))

    def check_seq_classifier_loss(self, config, input_values, *args):
        model = HubertForSequenceClassification(config=config)
        model.to(torch_device)

        
        model.eval()

        input_values = input_values[:3]
        attention_mask = torch.ones(input_values.shape, device=torch_device, dtype=torch.long)

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        labels = ids_tensor((input_values.shape[0], 1), len(model.config.id2label))

        
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0
            attention_mask[i, input_lengths[i] :] = 0

        masked_loss = model(input_values, attention_mask=attention_mask, labels=labels).loss.item()
        unmasked_loss = model(input_values, labels=labels).loss.item()

        self.parent.assertTrue(isinstance(masked_loss, float))
        self.parent.assertTrue(isinstance(unmasked_loss, float))
        self.parent.assertTrue(masked_loss != unmasked_loss)

    def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = HubertForCTC(config=config)
        model.to(torch_device)
        model.train()

        
        model.freeze_feature_encoder()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                
                
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()

    def check_seq_classifier_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = HubertForSequenceClassification(config=config)
        model.to(torch_device)
        model.train()

        
        model.freeze_base_model()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        labels = ids_tensor((input_values.shape[0], 1), len(model.config.id2label))

        
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()

    def check_labels_out_of_vocab(self, config, input_values, *args):
        model = HubertForCTC(config)
        model.to(torch_device)
        model.train()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size + 100)

        with pytest.raises(ValueError):
            model(input_values, labels=labels)

    def prepare_config_and_inputs_for_common(self):
        config, input_values, attention_mask = self.prepare_config_and_inputs()
        inputs_dict = {"input_values": input_values, "attention_mask": attention_mask}
        return config, inputs_dict


@require_torch
class HubertModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (HubertForCTC, HubertForSequenceClassification, HubertModel) if is_torch_available() else ()
    pipeline_model_mapping = (
        {
            "audio-classification": HubertForSequenceClassification,
            "automatic-speech-recognition": HubertForCTC,
            "feature-extraction": HubertModel,
        }
        if is_torch_available()
        else {}
    )
    fx_compatible = True
    test_pruning = False
    test_headmasking = False

    def setUp(self):
        self.model_tester = HubertModelTester(self)
        self.config_tester = ConfigTester(self, config_class=HubertConfig, hidden_size=37)

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

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_ctc_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_loss(*config_and_inputs)

    def test_seq_classifier_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_loss(*config_and_inputs)

    def test_ctc_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_training(*config_and_inputs)

    def test_seq_classifier_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_training(*config_and_inputs)

    def test_labels_out_of_vocab(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_labels_out_of_vocab(*config_and_inputs)

    
    def test_inputs_embeds(self):
        pass

    
    def test_forward_signature(self):
        pass

    
    
    def test_resize_tokens_embeddings(self):
        pass

    
    
    
    def test_model_common_attributes(self):
        pass

    def test_retain_grad_hidden_states_attentions(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        
        model.config.layerdrop = 0.0

        input_values = inputs_dict["input_values"]

        input_lengths = torch.tensor(
            [input_values.shape[1] for _ in range(input_values.shape[0])], dtype=torch.long, device=torch_device
        )
        output_lengths = model._get_feat_extract_output_lengths(input_lengths)

        labels = ids_tensor((input_values.shape[0], output_lengths[0] - 2), self.model_tester.vocab_size)
        inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"])
        inputs_dict["labels"] = labels

        outputs = model(**inputs_dict)

        output = outputs[0]

        
        hidden_states = outputs.hidden_states[0]
        attentions = outputs.attentions[0]

        hidden_states.retain_grad()
        attentions.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(hidden_states.grad)
        self.assertIsNotNone(attentions.grad)

    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                uniform_init_parms = [
                    "conv.weight",
                    "conv.parametrizations.weight",
                    "masked_spec_embed",
                    "quantizer.weight_proj.weight",
                ]
                if param.requires_grad:
                    if any(x in name for x in uniform_init_parms):
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )
                    else:
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )

    
    def _create_and_check_torch_fx_tracing(self, config, inputs_dict, output_loss=False):
        
        self.skipTest("torch 2.1 breaks torch fx tests for wav2vec2/hubert.")

        if not is_torch_fx_available() or not self.fx_compatible:
            return

        configs_no_init = _config_zero_init(config)  
        configs_no_init.return_dict = False

        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, return_labels=output_loss)

            try:
                if model.config.is_encoder_decoder:
                    model.config.use_cache = False  
                    labels = inputs.get("labels", None)
                    input_names = [
                        "attention_mask",
                        "decoder_attention_mask",
                        "decoder_input_ids",
                        "input_features",
                        "input_ids",
                        "input_values",
                    ]
                    if labels is not None:
                        input_names.append("labels")

                    filtered_inputs = {k: v for (k, v) in inputs.items() if k in input_names}
                    input_names = list(filtered_inputs.keys())

                    model_output = model(**filtered_inputs)

                    traced_model = symbolic_trace(model, input_names)
                    traced_output = traced_model(**filtered_inputs)
                else:
                    input_names = [
                        "attention_mask",
                        "bbox",
                        "input_features",
                        "input_ids",
                        "input_values",
                        "pixel_values",
                        "token_type_ids",
                        "visual_feats",
                        "visual_pos",
                    ]

                    labels = inputs.get("labels", None)
                    start_positions = inputs.get("start_positions", None)
                    end_positions = inputs.get("end_positions", None)
                    if labels is not None:
                        input_names.append("labels")
                    if start_positions is not None:
                        input_names.append("start_positions")
                    if end_positions is not None:
                        input_names.append("end_positions")

                    filtered_inputs = {k: v for (k, v) in inputs.items() if k in input_names}
                    input_names = list(filtered_inputs.keys())

                    model_output = model(**filtered_inputs)

                    traced_model = symbolic_trace(model, input_names)
                    traced_output = traced_model(**filtered_inputs)

            except Exception as e:
                self.fail(f"Couldn't trace module: {e}")

            def flatten_output(output):
                flatten = []
                for x in output:
                    if isinstance(x, (tuple, list)):
                        flatten += flatten_output(x)
                    elif not isinstance(x, torch.Tensor):
                        continue
                    else:
                        flatten.append(x)
                return flatten

            model_output = flatten_output(model_output)
            traced_output = flatten_output(traced_output)
            num_outputs = len(model_output)

            for i in range(num_outputs):
                self.assertTrue(
                    torch.allclose(model_output[i], traced_output[i]),
                    f"traced {i}th output doesn't match model {i}th output for {model_class}",
                )

            
            with tempfile.TemporaryDirectory() as tmp_dir_name:
                pkl_file_name = os.path.join(tmp_dir_name, "model.pkl")
                try:
                    with open(pkl_file_name, "wb") as f:
                        pickle.dump(traced_model, f)
                    with open(pkl_file_name, "rb") as f:
                        loaded = pickle.load(f)
                except Exception as e:
                    self.fail(f"Couldn't serialize / deserialize the traced model: {e}")

                loaded_output = loaded(**filtered_inputs)
                loaded_output = flatten_output(loaded_output)

                for i in range(num_outputs):
                    self.assertTrue(
                        torch.allclose(model_output[i], loaded_output[i]),
                        f"serialized model {i}th output doesn't match model {i}th output for {model_class}",
                    )

    
    def _mock_init_weights(self, module):
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data.fill_(3)
        if hasattr(module, "weight_g") and module.weight_g is not None:
            module.weight_g.data.fill_(3)
        if hasattr(module, "weight_v") and module.weight_v is not None:
            module.weight_v.data.fill_(3)
        if hasattr(module, "bias") and module.bias is not None:
            module.bias.data.fill_(3)
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)

    @unittest.skip(reason="Feed forward chunking is not implemented")
    def test_feed_forward_chunking(self):
        pass

    @slow
    def test_model_from_pretrained(self):
        model = HubertModel.from_pretrained("facebook/hubert-base-ls960")
        self.assertIsNotNone(model)


@require_torch
class HubertRobustModelTest(ModelTesterMixin, unittest.TestCase):
    all_model_classes = (HubertForCTC, HubertForSequenceClassification, HubertModel) if is_torch_available() else ()
    test_pruning = False
    test_headmasking = False

    def setUp(self):
        self.model_tester = HubertModelTester(
            self, conv_stride=(3, 3, 3), feat_extract_norm="layer", do_stable_layer_norm=True
        )
        self.config_tester = ConfigTester(self, config_class=HubertConfig, hidden_size=37)

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

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_batched_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_batch_inference(*config_and_inputs)

    def test_ctc_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_loss(*config_and_inputs)

    def test_seq_classifier_loss_inference(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_loss(*config_and_inputs)

    def test_ctc_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_ctc_training(*config_and_inputs)

    def test_seq_classifier_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_seq_classifier_training(*config_and_inputs)

    def test_labels_out_of_vocab(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_labels_out_of_vocab(*config_and_inputs)

    
    def test_inputs_embeds(self):
        pass

    
    def test_forward_signature(self):
        pass

    
    
    def test_resize_tokens_embeddings(self):
        pass

    
    
    
    def test_model_common_attributes(self):
        pass

    def test_retain_grad_hidden_states_attentions(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.output_hidden_states = True
        config.output_attentions = True

        
        model_class = self.all_model_classes[0]
        model = model_class(config)
        model.to(torch_device)

        
        model.config.layerdrop = 0.0

        input_values = inputs_dict["input_values"]

        input_lengths = torch.tensor(
            [input_values.shape[1] for _ in range(input_values.shape[0])], dtype=torch.long, device=torch_device
        )
        output_lengths = model._get_feat_extract_output_lengths(input_lengths)

        labels = ids_tensor((input_values.shape[0], output_lengths[0] - 2), self.model_tester.vocab_size)
        inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"])
        inputs_dict["labels"] = labels

        outputs = model(**inputs_dict)

        output = outputs[0]

        
        hidden_states = outputs.hidden_states[0]
        attentions = outputs.attentions[0]

        hidden_states.retain_grad()
        attentions.retain_grad()

        output.flatten()[0].backward(retain_graph=True)

        self.assertIsNotNone(hidden_states.grad)
        self.assertIsNotNone(attentions.grad)

    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                uniform_init_parms = [
                    "conv.weight",
                    "conv.parametrizations.weight",
                    "masked_spec_embed",
                    "quantizer.weight_proj.weight",
                ]
                if param.requires_grad:
                    if any(x in name for x in uniform_init_parms):
                        self.assertTrue(
                            -1.0 <= ((param.data.mean() * 1e9).round() / 1e9).item() <= 1.0,
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )
                    else:
                        self.assertIn(
                            ((param.data.mean() * 1e9).round() / 1e9).item(),
                            [0.0, 1.0],
                            msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                        )

    
    def _mock_init_weights(self, module):
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data.fill_(3)
        if hasattr(module, "weight_g") and module.weight_g is not None:
            module.weight_g.data.fill_(3)
        if hasattr(module, "weight_v") and module.weight_v is not None:
            module.weight_v.data.fill_(3)
        if hasattr(module, "bias") and module.bias is not None:
            module.bias.data.fill_(3)
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)

    @unittest.skip(reason="Feed forward chunking is not implemented")
    def test_feed_forward_chunking(self):
        pass

    @slow
    def test_model_from_pretrained(self):
        model = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")
        self.assertIsNotNone(model)


@require_torch
class HubertUtilsTest(unittest.TestCase):
    def test_compute_mask_indices(self):
        batch_size = 4
        sequence_length = 60
        mask_prob = 0.5
        mask_length = 1

        mask = _compute_mask_indices((batch_size, sequence_length), mask_prob, mask_length)
        mask = torch.from_numpy(mask).to(torch_device)

        self.assertListEqual(mask.sum(axis=-1).tolist(), [mask_prob * sequence_length for _ in range(batch_size)])

    def test_compute_mask_indices_overlap(self):
        batch_size = 4
        sequence_length = 80
        mask_prob = 0.5
        mask_length = 4

        mask = _compute_mask_indices((batch_size, sequence_length), mask_prob, mask_length)
        mask = torch.from_numpy(mask).to(torch_device)

        
        for batch_sum in mask.sum(axis=-1):
            self.assertTrue(int(batch_sum) <= mask_prob * sequence_length)


@require_torch
@require_soundfile
@slow
class HubertModelIntegrationTest(unittest.TestCase):
    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").filter(
            lambda x: x["id"] in [f"1272-141231-000{i}" for i in range(num_samples)]
        )[:num_samples]["audio"]

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

    def _load_superb(self, task, num_samples):
        from datasets import load_dataset

        ds = load_dataset("anton-l/superb_dummy", task, split="test")

        return ds[:num_samples]

    def test_inference_ctc_batched(self):
        model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft", torch_dtype=torch.float16).to(
            torch_device
        )
        processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft", do_lower_case=True)

        input_speech = self._load_datasamples(2)

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

        input_values = inputs.input_values.half().to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)

        with torch.no_grad():
            logits = model(input_values, attention_mask=attention_mask).logits

        predicted_ids = torch.argmax(logits, dim=-1)
        predicted_trans = processor.batch_decode(predicted_ids)

        EXPECTED_TRANSCRIPTIONS = [
            "a man said to the universe sir i exist",
            "sweat covered brion's body trickling into the tight loin cloth that was the only garment he wore",
        ]
        self.assertListEqual(predicted_trans, EXPECTED_TRANSCRIPTIONS)

    def test_inference_keyword_spotting(self):
        model = HubertForSequenceClassification.from_pretrained(
            "superb/hubert-base-superb-ks", torch_dtype=torch.float16
        ).to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/hubert-base-superb-ks")
        input_data = self._load_superb("ks", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.half().to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        predicted_logits, predicted_ids = torch.max(outputs.logits, dim=-1)

        expected_labels = [2, 6, 10, 9]
        
        expected_logits = torch.tensor([7.6692, 17.7795, 11.1562, 11.8232], dtype=torch.float16, device=torch_device)

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=3e-2))

    def test_inference_intent_classification(self):
        model = HubertForSequenceClassification.from_pretrained(
            "superb/hubert-base-superb-ic", torch_dtype=torch.float16
        ).to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/hubert-base-superb-ic")
        input_data = self._load_superb("ic", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.half().to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)

        predicted_logits_action, predicted_ids_action = torch.max(outputs.logits[:, :6], dim=-1)
        predicted_logits_object, predicted_ids_object = torch.max(outputs.logits[:, 6:20], dim=-1)
        predicted_logits_location, predicted_ids_location = torch.max(outputs.logits[:, 20:24], dim=-1)

        expected_labels_action = [1, 0, 4, 3]
        expected_logits_action = torch.tensor(
            [5.9052, 12.5865, 4.4840, 10.0240], dtype=torch.float16, device=torch_device
        )
        expected_labels_object = [1, 10, 3, 4]
        expected_logits_object = torch.tensor(
            [5.5316, 11.7946, 8.1672, 23.2415], dtype=torch.float16, device=torch_device
        )
        expected_labels_location = [0, 0, 0, 1]
        expected_logits_location = torch.tensor(
            [5.2053, 8.9577, 10.0447, 8.1481], dtype=torch.float16, device=torch_device
        )

        self.assertListEqual(predicted_ids_action.tolist(), expected_labels_action)
        self.assertListEqual(predicted_ids_object.tolist(), expected_labels_object)
        self.assertListEqual(predicted_ids_location.tolist(), expected_labels_location)

        
        self.assertTrue(torch.allclose(predicted_logits_action, expected_logits_action, atol=3e-1))
        self.assertTrue(torch.allclose(predicted_logits_object, expected_logits_object, atol=3e-1))
        self.assertTrue(torch.allclose(predicted_logits_location, expected_logits_location, atol=3e-1))

    def test_inference_speaker_identification(self):
        model = HubertForSequenceClassification.from_pretrained(
            "superb/hubert-base-superb-sid", torch_dtype=torch.float16
        ).to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/hubert-base-superb-sid")
        input_data = self._load_superb("si", 4)

        output_logits = []
        with torch.no_grad():
            for example in input_data["speech"]:
                input = processor(example, return_tensors="pt", padding=True)
                output = model(input.input_values.half().to(torch_device), attention_mask=None)
                output_logits.append(output.logits[0])
        output_logits = torch.stack(output_logits)
        predicted_logits, predicted_ids = torch.max(output_logits, dim=-1)

        expected_labels = [5, 1, 1, 3]
        
        expected_logits = torch.tensor(
            [78231.5547, 123166.6094, 122785.4141, 84851.2969], dtype=torch.float16, device=torch_device
        )

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=10))

    def test_inference_emotion_recognition(self):
        model = HubertForSequenceClassification.from_pretrained(
            "superb/hubert-base-superb-er", torch_dtype=torch.float16
        ).to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("superb/hubert-base-superb-er")
        input_data = self._load_superb("er", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)

        input_values = inputs.input_values.half().to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        predicted_logits, predicted_ids = torch.max(outputs.logits, dim=-1)

        expected_labels = [1, 1, 2, 2]
        
        expected_logits = torch.tensor([2.8384, 2.3389, 3.8564, 4.5558], dtype=torch.float16, device=torch_device)

        self.assertListEqual(predicted_ids.tolist(), expected_labels)
        
        self.assertTrue(torch.allclose(predicted_logits, expected_logits, atol=1e-1))

    def test_inference_distilhubert(self):
        model = HubertModel.from_pretrained("ntu-spml/distilhubert").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("ntu-spml/distilhubert")

        
        input_speech = self._load_datasamples(1)

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

        input_values = inputs.input_values.to(torch_device)

        with torch.no_grad():
            outputs = model(input_values).last_hidden_state

        
        expected_outputs_first = torch.tensor(
            [
                [
                    [-0.3505, 0.1167, 0.0608, 0.1294],
                    [-0.3085, 0.0481, 0.1106, 0.0955],
                    [-0.3107, -0.0391, 0.0739, 0.1360],
                    [-0.2385, -0.1795, -0.0928, 0.2389],
                ]
            ],
            device=torch_device,
        )
        expected_outputs_last = torch.tensor(
            [
                [
                    [-0.0732, 0.0255, 0.0529, -0.1372],
                    [-0.0812, 0.1259, 0.0564, -0.0438],
                    [-0.0054, 0.0758, -0.0002, -0.1617],
                    [0.0133, -0.0320, -0.0687, 0.0062],
                ]
            ],
            device=torch_device,
        )
        expected_output_sum = -3776.0730

        self.assertTrue(torch.allclose(outputs[:, :4, :4], expected_outputs_first, atol=5e-3))
        self.assertTrue(torch.allclose(outputs[:, -4:, -4:], expected_outputs_last, atol=5e-3))
        self.assertTrue(abs(outputs.sum() - expected_output_sum) < 0.1)