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 unispeechsat model import math import unittest import numpy as np import pytest from datasets import loaddataset from transformers import unispeechsatconfig istorchavailable from transformers testingutils import requiresoundfile requiretorch slow torchdevice from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin configzeroinit floatstensor idstensor randomattentionmask from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import unispeechsatforaudioframeclassification unispeechsatforctc unispeechsatforpretraining unispeechsatforsequenceclassification unispeechsatforxvector unispeechsatmodel wav2vec2featureextractor wav2vec2processor class unispeechsatmodeltester 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 masktimeprob0 5 masktimelength2 vocabsize32 dostablelayernormfalse tdnndim32 32 tdnnkernel3 3 tdnndilation1 1 xvectoroutputdim32 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 masktimeprob masktimeprob self masktimelength masktimelength self tdnndim tdnndim self tdnnkernel tdnnkernel self tdnndilation tdnndilation self xvectoroutputdim xvectoroutputdim 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 unispeechsatconfig hiddensizeself hiddensize featextractnormself featextractnorm featextractdropoutself featextractdropout featextractactivationself featextractactivation convdimself convdim convstrideself convstride convkernelself convkernel convbiasself convbias numconvposembeddingsself numconvposembeddings numconvposembeddinggroupsself numconvposembeddinggroups masktimeprobself masktimeprob masktimelengthself masktimelength numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads hiddendropoutprobself hiddendropoutprob intermediatesizeself intermediatesize layernormepsself layernormeps hiddenactself hiddenact initializerrangeself initializerrange vocabsizeself vocabsize tdnndimself tdnndim tdnnkernelself tdnnkernel tdnndilationself tdnndilation xvectoroutputdimself xvectoroutputdim def createandcheckmodelself config inputvalues attentionmask model unispeechsatmodelconfigconfig 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 unispeechsatmodelconfigconfig 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 unispeechsatforctcconfigconfig 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 unispeechsatforsequenceclassificationconfigconfig 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 unispeechsatforctcconfigconfig 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 unispeechsatforsequenceclassificationconfigconfig 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 checkxvectortrainingself config args config ctczeroinfinity true model unispeechsatforxvectorconfigconfig model totorchdevice model train freeze everything but the classification head model freezebasemodel use a longer sequence length to account for tdnn temporal downsampling inputvalues floatstensorself batchsize self seqlength 2 scale1 0 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 unispeechsatforctcconfig 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 unispeechsatmodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses unispeechsatforctc unispeechsatforpretraining unispeechsatmodel unispeechsatforsequenceclassification unispeechsatforaudioframeclassification unispeechsatforxvector if istorchavailable else pipelinemodelmapping audioclassification unispeechsatforsequenceclassification automaticspeechrecognition unispeechsatforctc featureextraction unispeechsatmodel if istorchavailable else testpruning false testheadmasking false testtorchscript false def setupself self modeltester unispeechsatmodeltesterself self configtester configtesterself configclassunispeechsatconfig 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 testxvectortrainself configandinputs self modeltester prepareconfigandinputs self modeltester checkxvectortrainingconfigandinputs def testlabelsoutofvocabself configandinputs self modeltester prepareconfigandinputs self modeltester checklabelsoutofvocabconfigandinputs unispeechsat has no inputsembeds def testinputsembedsself pass inputids is renamed to inputvalues def testforwardsignatureself pass unispeechsat cannot resize token embeddings since it has no tokens embeddings def testresizetokensembeddingsself pass unispeechsat 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 codevectors quantizer weightproj weight projecthid weight projecthid bias projectq weight projectq bias featureprojection projection weight featureprojection projection bias labelembeddingsconcat objective 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 codevectors and module codevectors is not none module codevectors data fill3 if hasattrmodule maskedspecembed and module maskedspecembed is not none module maskedspecembed data fill3 def testmaskfeatureprobctcself model unispeechsatforctc frompretrained hfinternaltestingtinyrandomunispeechsat maskfeatureprob0 2 maskfeaturelength2 model totorchdevice train processor wav2vec2processor frompretrained hfinternaltestingtinyrandomunispeechsat returnattentionmasktrue batchdurationinseconds 1 3 2 6 inputfeatures np random random16000 s for s in batchdurationinseconds batch processor inputfeatures paddingtrue samplingrateprocessor featureextractor samplingrate returntensorspt logits model inputvaluesbatchinputvalues totorchdevice attentionmaskbatchattentionmask totorchdevice logits self assertequallogits shape 4 1498 32 def testmasktimeprobctcself model unispeechsatforctc frompretrained hfinternaltestingtinyrandomunispeechsat masktimeprob0 2 masktimelength2 model totorchdevice train processor wav2vec2processor frompretrained hfinternaltestingtinyrandomunispeechsat returnattentionmasktrue batchdurationinseconds 1 3 2 6 inputfeatures np random random16000 s for s in batchdurationinseconds batch processor inputfeatures paddingtrue samplingrateprocessor featureextractor samplingrate returntensorspt logits model inputvaluesbatchinputvalues totorchdevice attentionmaskbatchattentionmask totorchdevice logits self assertequallogits shape 4 1498 32 unittest skipreasonfeed forward chunking is not implemented def testfeedforwardchunkingself pass slow def testmodelfrompretrainedself model unispeechsatmodel frompretrainedmicrosoftunispeechsatbaseplus self assertisnotnonemodel requiretorch class unispeechsatrobustmodeltestmodeltestermixin unittest testcase allmodelclasses unispeechsatforctc unispeechsatforpretraining unispeechsatmodel unispeechsatforsequenceclassification if istorchavailable else testpruning false testheadmasking false testtorchscript false def setupself self modeltester unispeechsatmodeltester self convstride3 3 3 featextractnormlayer dostablelayernormtrue self configtester configtesterself configclassunispeechsatconfig 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 unispeechsat has no inputsembeds def testinputsembedsself pass inputids is renamed to inputvalues def testforwardsignatureself pass unispeechsat cannot resize token embeddings since it has no tokens embeddings def testresizetokensembeddingsself pass unispeechsat 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 codevectors quantizer weightproj weight projecthid weight projecthid bias projectq weight projectq bias featureprojection projection weight featureprojection projection bias labelembeddingsconcat objective 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 codevectors and module codevectors is not none module codevectors data fill3 if hasattrmodule maskedspecembed and module maskedspecembed is not none module maskedspecembed data fill3 def testmaskfeatureprobctcself model unispeechsatforctc frompretrained hfinternaltestingtinyrandomunispeechsat maskfeatureprob0 2 maskfeaturelength2 model totorchdevice train processor wav2vec2processor frompretrained hfinternaltestingtinyrandomunispeechsat returnattentionmasktrue batchdurationinseconds 1 3 2 6 inputfeatures np random random16000 s for s in batchdurationinseconds batch processor inputfeatures paddingtrue samplingrateprocessor featureextractor samplingrate returntensorspt logits model inputvaluesbatchinputvalues totorchdevice attentionmaskbatchattentionmask totorchdevice logits self assertequallogits shape 4 1498 32 def testmasktimeprobctcself model unispeechsatforctc frompretrained hfinternaltestingtinyrandomunispeechsat masktimeprob0 2 masktimelength2 model totorchdevice train processor wav2vec2processor frompretrained hfinternaltestingtinyrandomunispeechsat returnattentionmasktrue batchdurationinseconds 1 3 2 6 inputfeatures np random random16000 s for s in batchdurationinseconds batch processor inputfeatures paddingtrue samplingrateprocessor featureextractor samplingrate returntensorspt logits model inputvaluesbatchinputvalues totorchdevice attentionmaskbatchattentionmask totorchdevice logits self assertequallogits shape 4 1498 32 def testmasktimefeatureprobctcsinglebatchself model unispeechsatforctc frompretrained hfinternaltestingtinyrandomunispeechsat masktimeprob0 2 maskfeatureprob0 2 masktimelength2 maskfeaturelength2 model totorchdevice train processor wav2vec2processor frompretrained hfinternaltestingtinyrandomunispeechsat returnattentionmasktrue batchdurationinseconds 6 inputfeatures np random random16000 s for s in batchdurationinseconds batch processor inputfeatures paddingtrue samplingrateprocessor featureextractor samplingrate returntensorspt logits model inputvaluesbatchinputvalues totorchdevice attentionmaskbatchattentionmask totorchdevice logits self assertequallogits shape 1 1498 32 unittest skipreasonfeed forward chunking is not implemented def testfeedforwardchunkingself pass slow def testmodelfrompretrainedself model unispeechsatmodel frompretrainedmicrosoftunispeechsatlarge self assertisnotnonemodel requiretorch requiresoundfile slow class unispeechsatmodelintegrationtestunittest testcase def loaddatasamplesself numsamples 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 ds loaddatasetantonlsuperbdummy task splittest return ds numsamples def testinferenceencoderbaseself model unispeechsatmodel frompretrainedmicrosoftunispeechsatbaseplus model totorchdevice featureextractor wav2vec2featureextractor frompretrained facebookwav2vec2base returnattentionmasktrue inputspeech self loaddatasamples2 inputsdict featureextractorinputspeech returntensorspt paddingtrue with torch nograd outputs model inputsdict inputvalues totorchdevice attentionmaskinputsdict attentionmask totorchdevice fmt off expectedhiddenstatesslice torch tensor 0 0743 0 1384 0 0845 0 1704 0 0954 0 1936 0 1123 0 2095 devicetorchdevice fmt on self asserttruetorch allcloseoutputs lasthiddenstate 2 2 expectedhiddenstatesslice atol1e3 def testinferenceencoderlargeself model unispeechsatmodel frompretrainedmicrosoftunispeechsatlarge model totorchdevice featureextractor wav2vec2featureextractor frompretrainedfacebookwav2vec2largexlsr53 inputspeech self loaddatasamples2 inputsdict featureextractorinputspeech returntensorspt paddingtrue with torch nograd outputs model inputsdict inputvalues totorchdevice attentionmaskinputsdict attentionmask totorchdevice fmt off expectedhiddenstatesslice torch tensor 0 1172 0 0797 0 0012 0 0213 0 1225 0 1277 0 0668 0 0585 devicetorchdevice fmt on self asserttruetorch allcloseoutputs lasthiddenstate 2 2 expectedhiddenstatesslice atol1e3 def testinferencediarizationself model unispeechsatforaudioframeclassification frompretrainedmicrosoftunispeechsatbaseplussd to torchdevice processor wav2vec2featureextractor frompretrainedmicrosoftunispeechsatbaseplussd inputdata self loadsuperbsd 4 inputs processorinputdataspeech returntensorspt paddingtrue samplingrate16000 inputvalues inputs inputvalues totorchdevice attentionmask inputs attentionmask totorchdevice with torch nograd outputs modelinputvalues attentionmaskattentionmask labels is a onehot array of shape numframes numspeakers labels outputs logits 0 long s3prl logits for the same batch expectedlogits torch tensor 5 6119 5 5845 3 7772 5 4824 3 6914 5 1619 4 7560 5 0496 6 3785 4 8365 5 5863 5 4149 5 5639 4 8469 6 1511 4 0052 6 0355 3 7414 5 5968 4 8061 5 4620 4 7310 5 5864 4 6078 5 9493 4 8963 4 4050 5 4476 4 1755 5 1395 4 0272 4 3705 devicetorchdevice self assertequallabels0 0 sum 270 self assertequallabels0 1 sum 647 todo update the tolerance after the ci moves to torch 1 10 self asserttruetorch allcloseoutputs logits 4 expectedlogits atol1e2 def testinferencespeakerverificationself model unispeechsatforxvector frompretrainedmicrosoftunispeechsatbaseplussv totorchdevice processor wav2vec2featureextractor frompretrainedmicrosoftunispeechsatbaseplussv inputdata self loadsuperbsi 4 inputs processorinputdataspeech returntensorspt paddingtrue labels torch tensor5 1 1 3 devicetorchdevice t with torch nograd inputvalues inputs inputvalues totorchdevice attentionmask inputs attentionmask totorchdevice outputs modelinputvalues attentionmaskattentionmask labelslabels embeddings torch nn functional normalizeoutputs embeddings dim1 cosinesim torch nn cosinesimilaritydim1 id10002 vs id10002 self assertalmostequalcosinesimembeddings1 embeddings2 item 0 9671 3 id10006 vs id10002 self assertalmostequalcosinesimembeddings0 embeddings1 item 0 4941 3 id10002 vs id10004 self assertalmostequalcosinesimembeddings2 embeddings3 item 0 5616 3 todo update the tolerance after the ci moves to torch 1 10 self assertalmostequaloutputs loss item 18 5925 2 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 unispeechsat 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 freeze everything but the classification head use a longer sequence length to account for tdnn temporal downsampling pad input unispeechsat has no inputs_embeds input_ids is renamed to input_values unispeechsat cannot resize token embeddings since it has no tokens embeddings unispeechsat 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 unispeechsat has no inputs_embeds input_ids is renamed to input_values unispeechsat cannot resize token embeddings since it has no tokens embeddings unispeechsat 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 automatic decoding with librispeech fmt off fmt on fmt off fmt on labels is a one hot array of shape num_frames num_speakers s3prl logits for the same batch todo update the tolerance after the ci moves to torch 1 10 id10002 vs id10002 id10006 vs id10002 id10002 vs id10004 todo update the tolerance after the ci moves to torch 1 10
import math
import unittest

import numpy as np
import pytest
from datasets import load_dataset

from transformers import UniSpeechSatConfig, is_torch_available
from transformers.testing_utils import require_soundfile, require_torch, slow, torch_device

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 (
        UniSpeechSatForAudioFrameClassification,
        UniSpeechSatForCTC,
        UniSpeechSatForPreTraining,
        UniSpeechSatForSequenceClassification,
        UniSpeechSatForXVector,
        UniSpeechSatModel,
        Wav2Vec2FeatureExtractor,
        Wav2Vec2Processor,
    )


class UniSpeechSatModelTester:
    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,
        mask_time_prob=0.5,
        mask_time_length=2,
        vocab_size=32,
        do_stable_layer_norm=False,
        tdnn_dim=(32, 32),
        tdnn_kernel=(3, 3),
        tdnn_dilation=(1, 1),
        xvector_output_dim=32,
        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.mask_time_prob = mask_time_prob
        self.mask_time_length = mask_time_length
        self.tdnn_dim = tdnn_dim
        self.tdnn_kernel = tdnn_kernel
        self.tdnn_dilation = tdnn_dilation
        self.xvector_output_dim = xvector_output_dim
        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 UniSpeechSatConfig(
            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,
            mask_time_prob=self.mask_time_prob,
            mask_time_length=self.mask_time_length,
            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,
            tdnn_dim=self.tdnn_dim,
            tdnn_kernel=self.tdnn_kernel,
            tdnn_dilation=self.tdnn_dilation,
            xvector_output_dim=self.xvector_output_dim,
        )

    def create_and_check_model(self, config, input_values, attention_mask):
        model = UniSpeechSatModel(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 = UniSpeechSatModel(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 = UniSpeechSatForCTC(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 = UniSpeechSatForSequenceClassification(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 = UniSpeechSatForCTC(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 = UniSpeechSatForSequenceClassification(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_xvector_training(self, config, *args):
        config.ctc_zero_infinity = True
        model = UniSpeechSatForXVector(config=config)
        model.to(torch_device)
        model.train()

        
        model.freeze_base_model()

        
        input_values = floats_tensor([self.batch_size, self.seq_length * 2], scale=1.0)

        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 = UniSpeechSatForCTC(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 UniSpeechSatModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (
            UniSpeechSatForCTC,
            UniSpeechSatForPreTraining,
            UniSpeechSatModel,
            UniSpeechSatForSequenceClassification,
            UniSpeechSatForAudioFrameClassification,
            UniSpeechSatForXVector,
        )
        if is_torch_available()
        else ()
    )
    pipeline_model_mapping = (
        {
            "audio-classification": UniSpeechSatForSequenceClassification,
            "automatic-speech-recognition": UniSpeechSatForCTC,
            "feature-extraction": UniSpeechSatModel,
        }
        if is_torch_available()
        else {}
    )
    test_pruning = False
    test_headmasking = False
    test_torchscript = False

    def setUp(self):
        self.model_tester = UniSpeechSatModelTester(self)
        self.config_tester = ConfigTester(self, config_class=UniSpeechSatConfig, 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_xvector_train(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_xvector_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",
                    "codevectors",
                    "quantizer.weight_proj.weight",
                    "project_hid.weight",
                    "project_hid.bias",
                    "project_q.weight",
                    "project_q.bias",
                    "feature_projection.projection.weight",
                    "feature_projection.projection.bias",
                    "label_embeddings_concat",
                    "objective.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, "codevectors") and module.codevectors is not None:
            module.codevectors.data.fill_(3)
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)

    def test_mask_feature_prob_ctc(self):
        model = UniSpeechSatForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", mask_feature_prob=0.2, mask_feature_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )
        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

    def test_mask_time_prob_ctc(self):
        model = UniSpeechSatForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", mask_time_prob=0.2, mask_time_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

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

    @slow
    def test_model_from_pretrained(self):
        model = UniSpeechSatModel.from_pretrained("microsoft/unispeech-sat-base-plus")
        self.assertIsNotNone(model)


@require_torch
class UniSpeechSatRobustModelTest(ModelTesterMixin, unittest.TestCase):
    all_model_classes = (
        (UniSpeechSatForCTC, UniSpeechSatForPreTraining, UniSpeechSatModel, UniSpeechSatForSequenceClassification)
        if is_torch_available()
        else ()
    )
    test_pruning = False
    test_headmasking = False
    test_torchscript = False

    def setUp(self):
        self.model_tester = UniSpeechSatModelTester(
            self, conv_stride=(3, 3, 3), feat_extract_norm="layer", do_stable_layer_norm=True
        )
        self.config_tester = ConfigTester(self, config_class=UniSpeechSatConfig, 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",
                    "codevectors",
                    "quantizer.weight_proj.weight",
                    "project_hid.weight",
                    "project_hid.bias",
                    "project_q.weight",
                    "project_q.bias",
                    "feature_projection.projection.weight",
                    "feature_projection.projection.bias",
                    "label_embeddings_concat",
                    "objective.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, "codevectors") and module.codevectors is not None:
            module.codevectors.data.fill_(3)
        if hasattr(module, "masked_spec_embed") and module.masked_spec_embed is not None:
            module.masked_spec_embed.data.fill_(3)

    def test_mask_feature_prob_ctc(self):
        model = UniSpeechSatForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", mask_feature_prob=0.2, mask_feature_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

    def test_mask_time_prob_ctc(self):
        model = UniSpeechSatForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", mask_time_prob=0.2, mask_time_length=2
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", return_attention_mask=True
        )

        batch_duration_in_seconds = [1, 3, 2, 6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (4, 1498, 32))

    def test_mask_time_feature_prob_ctc_single_batch(self):
        model = UniSpeechSatForCTC.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat",
            mask_time_prob=0.2,
            mask_feature_prob=0.2,
            mask_time_length=2,
            mask_feature_length=2,
        )
        model.to(torch_device).train()
        processor = Wav2Vec2Processor.from_pretrained(
            "hf-internal-testing/tiny-random-unispeech-sat", return_attention_mask=True
        )

        batch_duration_in_seconds = [6]
        input_features = [np.random.random(16_000 * s) for s in batch_duration_in_seconds]

        batch = processor(
            input_features, padding=True, sampling_rate=processor.feature_extractor.sampling_rate, return_tensors="pt"
        )

        logits = model(
            input_values=batch["input_values"].to(torch_device),
            attention_mask=batch["attention_mask"].to(torch_device),
        ).logits

        self.assertEqual(logits.shape, (1, 1498, 32))

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

    @slow
    def test_model_from_pretrained(self):
        model = UniSpeechSatModel.from_pretrained("microsoft/unispeech-sat-large")
        self.assertIsNotNone(model)


@require_torch
@require_soundfile
@slow
class UniSpeechSatModelIntegrationTest(unittest.TestCase):
    def _load_datasamples(self, num_samples):
        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):
        ds = load_dataset("anton-l/superb_dummy", task, split="test")

        return ds[:num_samples]

    def test_inference_encoder_base(self):
        model = UniSpeechSatModel.from_pretrained("microsoft/unispeech-sat-base-plus")
        model.to(torch_device)
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
            "facebook/wav2vec2-base", return_attention_mask=True
        )
        input_speech = self._load_datasamples(2)

        inputs_dict = feature_extractor(input_speech, return_tensors="pt", padding=True)

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

        
        expected_hidden_states_slice = torch.tensor(
            [[[-0.0743, 0.1384],
              [-0.0845, 0.1704]],
             [[-0.0954, 0.1936],
              [-0.1123, 0.2095]]],
            device=torch_device,
        )
        

        self.assertTrue(torch.allclose(outputs.last_hidden_state[:, :2, -2:], expected_hidden_states_slice, atol=1e-3))

    def test_inference_encoder_large(self):
        model = UniSpeechSatModel.from_pretrained("microsoft/unispeech-sat-large")
        model.to(torch_device)
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-large-xlsr-53")
        input_speech = self._load_datasamples(2)

        inputs_dict = feature_extractor(input_speech, return_tensors="pt", padding=True)

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

        
        expected_hidden_states_slice = torch.tensor(
            [[[-0.1172, -0.0797],
              [-0.0012, 0.0213]],
             [[-0.1225, -0.1277],
              [-0.0668, -0.0585]]],
            device=torch_device,
        )
        

        self.assertTrue(torch.allclose(outputs.last_hidden_state[:, :2, -2:], expected_hidden_states_slice, atol=1e-3))

    def test_inference_diarization(self):
        model = UniSpeechSatForAudioFrameClassification.from_pretrained("microsoft/unispeech-sat-base-plus-sd").to(
            torch_device
        )
        processor = Wav2Vec2FeatureExtractor.from_pretrained("microsoft/unispeech-sat-base-plus-sd")
        input_data = self._load_superb("sd", 4)
        inputs = processor(input_data["speech"], return_tensors="pt", padding=True, sampling_rate=16_000)

        input_values = inputs.input_values.to(torch_device)
        attention_mask = inputs.attention_mask.to(torch_device)
        with torch.no_grad():
            outputs = model(input_values, attention_mask=attention_mask)
        
        labels = (outputs.logits > 0).long()

        
        expected_logits = torch.tensor(
            [
                [[-5.6119, -5.5845], [-3.7772, -5.4824], [-3.6914, -5.1619], [-4.7560, -5.0496]],
                [[-6.3785, -4.8365], [-5.5863, -5.4149], [-5.5639, -4.8469], [-6.1511, -4.0052]],
                [[-6.0355, -3.7414], [-5.5968, -4.8061], [-5.4620, -4.7310], [-5.5864, -4.6078]],
                [[-5.9493, -4.8963], [-4.4050, -5.4476], [-4.1755, -5.1395], [-4.0272, -4.3705]],
            ],
            device=torch_device,
        )
        self.assertEqual(labels[0, :, 0].sum(), 270)
        self.assertEqual(labels[0, :, 1].sum(), 647)
        
        self.assertTrue(torch.allclose(outputs.logits[:, :4], expected_logits, atol=1e-2))

    def test_inference_speaker_verification(self):
        model = UniSpeechSatForXVector.from_pretrained("microsoft/unispeech-sat-base-plus-sv").to(torch_device)
        processor = Wav2Vec2FeatureExtractor.from_pretrained("microsoft/unispeech-sat-base-plus-sv")
        input_data = self._load_superb("si", 4)

        inputs = processor(input_data["speech"], return_tensors="pt", padding=True)
        labels = torch.tensor([5, 1, 1, 3], device=torch_device).T

        with torch.no_grad():
            input_values = inputs.input_values.to(torch_device)
            attention_mask = inputs.attention_mask.to(torch_device)
            outputs = model(input_values, attention_mask=attention_mask, labels=labels)
        embeddings = torch.nn.functional.normalize(outputs.embeddings, dim=-1)

        cosine_sim = torch.nn.CosineSimilarity(dim=-1)
        
        self.assertAlmostEqual(cosine_sim(embeddings[1], embeddings[2]).item(), 0.9671, 3)
        
        self.assertAlmostEqual(cosine_sim(embeddings[0], embeddings[1]).item(), 0.4941, 3)
        
        self.assertAlmostEqual(cosine_sim(embeddings[2], embeddings[3]).item(), 0.5616, 3)

        
        self.assertAlmostEqual(outputs.loss.item(), 18.5925, 2)