codingutf8 2023 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 pop2piano model import copy import tempfile import unittest import numpy as np from datasets import loaddataset from transformers import pop2pianoconfig from transformers featureextractionutils import batchfeature from transformers testingutils import requireessentia requirelibrosa requireonnx requirescipy requiretorch slow torchdevice from transformers utils import isessentiaavailable islibrosaavailable isscipyavailable istorchavailable from generation testutils import generationtestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import pop2pianoforconditionalgeneration from transformers models pop2piano modelingpop2piano import pop2pianopretrainedmodelarchivelist from transformers pytorchutils import istorch180 else istorch180 false requiretorch class pop2pianomodeltester def init self parent vocabsize99 batchsize13 encoderseqlength7 decoderseqlength9 for common tests istrainingfalse useattentionmasktrue uselabelstrue hiddensize64 numhiddenlayers5 numattentionheads4 dff37 relativeattentionnumbuckets8 dropoutrate0 1 initializerfactor0 002 eostokenid1 padtokenid0 decoderstarttokenid0 scopenone decoderlayersnone self parent parent self batchsize batchsize self encoderseqlength encoderseqlength self decoderseqlength decoderseqlength for common tests self seqlength self decoderseqlength self istraining istraining self useattentionmask useattentionmask self uselabels uselabels self vocabsize vocabsize self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self dff dff self relativeattentionnumbuckets relativeattentionnumbuckets self dropoutrate dropoutrate self initializerfactor initializerfactor self eostokenid eostokenid self padtokenid padtokenid self decoderstarttokenid decoderstarttokenid self scope none self decoderlayers decoderlayers def prepareconfigandinputsself inputids idstensorself batchsize self encoderseqlength self vocabsize decoderinputids idstensorself batchsize self decoderseqlength self vocabsize attentionmask none decoderattentionmask none if self useattentionmask attentionmask idstensorself batchsize self encoderseqlength vocabsize2 decoderattentionmask idstensorself batchsize self decoderseqlength vocabsize2 lmlabels idstensorself batchsize self decoderseqlength self vocabsize if self uselabels else none return self getconfig inputids decoderinputids attentionmask decoderattentionmask lmlabels def getpipelineconfigself return pop2pianoconfig vocabsize166 pop2piano forces 100 extra tokens dmodelself hiddensize dffself dff dkvself hiddensize self numattentionheads numlayersself numhiddenlayers numdecoderlayersself decoderlayers numheadsself numattentionheads relativeattentionnumbucketsself relativeattentionnumbuckets dropoutrateself dropoutrate initializerfactorself initializerfactor eostokenidself eostokenid bostokenidself padtokenid padtokenidself padtokenid decoderstarttokenidself decoderstarttokenid def getconfigself return pop2pianoconfig vocabsizeself vocabsize dmodelself hiddensize dffself dff dkvself hiddensize self numattentionheads numlayersself numhiddenlayers numdecoderlayersself decoderlayers numheadsself numattentionheads relativeattentionnumbucketsself relativeattentionnumbuckets dropoutrateself dropoutrate initializerfactorself initializerfactor eostokenidself eostokenid bostokenidself padtokenid padtokenidself padtokenid decoderstarttokenidself decoderstarttokenid def checkpreparelmlabelsviashiftleft self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig model totorchdevice model eval make sure that lmlabels are correctly padded from the right lmlabels maskedfilllmlabels self decoderstarttokenid self eostokenid add causal pad token mask triangularmask torch trillmlabels newoneslmlabels shape logicalnot lmlabels maskedfilltriangularmask self padtokenid decoderinputids model shiftrightlmlabels for i decoderinputidsslice lmlabelsslice in enumeratezipdecoderinputids lmlabels first item self parent assertequaldecoderinputidsslice0 item self decoderstarttokenid if i decoderinputidsslice shape1 if i decoderinputids shape1 1 items before diagonal self parent assertlistequal decoderinputidsslice1 i 1 tolist lmlabelsslice i tolist pad items after diagonal if i decoderinputids shape1 2 self parent assertlistequal decoderinputidsslicei 2 tolist lmlabelsslicei 1 1 tolist else all items after square self parent assertlistequaldecoderinputidsslice1 tolist lmlabelsslice 1 tolist def createandcheckmodel self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig model totorchdevice model eval result model inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask result modelinputidsinputids decoderinputidsdecoderinputids decoderpast result pastkeyvalues encoderoutput result encoderlasthiddenstate self parent assertequalencoderoutput size self batchsize self encoderseqlength self hiddensize there should be numlayers key value embeddings stored in decoderpast self parent assertequallendecoderpast config numlayers there should be a self attn key a self attn value a cross attn key and a cross attn value stored in each decoderpast tuple self parent assertequallendecoderpast0 4 def createandcheckwithlmhead self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig totorchdevice eval outputs model inputidsinputids decoderinputidsdecoderinputids decoderattentionmaskdecoderattentionmask labelslmlabels self parent assertequallenoutputs 4 self parent assertequaloutputslogits size self batchsize self decoderseqlength self vocabsize self parent assertequaloutputsloss size def createandcheckdecodermodelpast self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig getdecoder totorchdevice eval first forward pass outputs modelinputids usecachetrue outputsusecacheconf modelinputids outputsnopast modelinputids usecachefalse self parent asserttruelenoutputs lenoutputsusecacheconf self parent asserttruelenoutputs lenoutputsnopast 1 output pastkeyvalues outputs totuple create hypothetical next token and extent to nextinputids nexttokens idstensorself batchsize 1 config vocabsize append to next inputids and nextinputids torch catinputids nexttokens dim1 outputfromnopast modelnextinputidslasthiddenstate outputfrompast modelnexttokens pastkeyvaluespastkeyvalueslasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 1 randomsliceidx detach outputfrompastslice outputfrompast 0 randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandcheckdecodermodelattentionmaskpast self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig getdecoder model totorchdevice model eval create attention mask attnmask torch onesinputids shape dtypetorch long devicetorchdevice halfseqlength inputids shape1 2 attnmask halfseqlength 0 first forward pass output pastkeyvalues modelinputids attentionmaskattnmask usecachetrue totuple create hypothetical next token and extent to nextinputids nexttokens idstensorself batchsize 1 config vocabsize change a random masked slice from inputids randomseqidxtochange idstensor1 halfseqlength item 1 randomothernexttokens idstensorself batchsize 1 config vocabsize squeeze1 inputids randomseqidxtochange randomothernexttokens append to next inputids and attnmask nextinputids torch catinputids nexttokens dim1 attnmask torch cat attnmask torch onesattnmask shape0 1 dtypetorch long devicetorchdevice dim1 get two different outputs outputfromnopast modelnextinputids attentionmaskattnmasklasthiddenstate outputfrompast modelnexttokens pastkeyvaluespastkeyvalues attentionmaskattnmask lasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 1 randomsliceidx detach outputfrompastslice outputfrompast 0 randomsliceidx detach test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandcheckdecodermodelpastlargeinputs self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig getdecoder totorchdevice eval first forward pass outputs modelinputids attentionmaskattentionmask usecachetrue output pastkeyvalues outputs totuple create hypothetical multiple next token and extent to nextinputids nexttokens idstensorself batchsize 3 config vocabsize nextmask idstensorself batchsize 3 vocabsize2 append to next inputids and nextinputids torch catinputids nexttokens dim1 nextattentionmask torch catattentionmask nextmask dim1 outputfromnopast modelnextinputids attentionmasknextattentionmasklasthiddenstate outputfrompast modelnexttokens attentionmasknextattentionmask pastkeyvaluespastkeyvalues lasthiddenstate select random slice randomsliceidx idstensor1 outputfrompast shape1 item outputfromnopastslice outputfromnopast 3 randomsliceidx detach outputfrompastslice outputfrompast randomsliceidx detach self parent asserttrueoutputfrompastslice shape1 nexttokens shape1 test that outputs are equal for slice self parent asserttruetorch allcloseoutputfrompastslice outputfromnopastslice atol1e3 def createandcheckgeneratewithpastkeyvalues self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig totorchdevice eval torch manualseed0 outputwithoutpastcache model generate inputids 1 numbeams2 maxlength5 dosampletrue usecachefalse torch manualseed0 outputwithpastcache model generateinputids 1 numbeams2 maxlength5 dosampletrue self parent asserttruetorch alloutputwithpastcache outputwithoutpastcache def createandcheckmodelfp16forward self config inputids decoderinputids attentionmask decoderattentionmask lmlabels model pop2pianoforconditionalgenerationconfigconfig totorchdevice half eval output modelinputids decoderinputidsinputids attentionmaskattentionmask encoderlasthiddenstate self parent assertfalsetorch isnanoutput any item def createandcheckencoderdecodersharedweights self config inputids decoderinputids attentionmask decoderattentionmask lmlabels for modelclass in pop2pianoforconditionalgeneration torch manualseed0 model modelclassconfigconfig totorchdevice eval load state dict copies weights but does not tie them model encoder loadstatedictmodel decoder statedict strictfalse torch manualseed0 tiedconfig copy deepcopyconfig tiedconfig tieencoderdecoder true tiedmodel modelclassconfigtiedconfig totorchdevice eval modelresult model inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask tiedmodelresult tiedmodel inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask check that models has less parameters self parent assertless sump numel for p in tiedmodel parameters sump numel for p in model parameters randomsliceidx idstensor1 modelresult0 shape1 item check that outputs are equal self parent asserttrue torch allclose modelresult00 randomsliceidx tiedmodelresult00 randomsliceidx atol1e4 check that outputs after saving and loading are equal with tempfile temporarydirectory as tmpdirname tiedmodel savepretrainedtmpdirname tiedmodel modelclass frompretrainedtmpdirname tiedmodel totorchdevice tiedmodel eval check that models has less parameters self parent assertless sump numel for p in tiedmodel parameters sump numel for p in model parameters randomsliceidx idstensor1 modelresult0 shape1 item tiedmodelresult tiedmodel inputidsinputids decoderinputidsdecoderinputids attentionmaskattentionmask decoderattentionmaskdecoderattentionmask check that outputs are equal self parent asserttrue torch allclose modelresult00 randomsliceidx tiedmodelresult00 randomsliceidx atol1e4 def checkresizeembeddingspop2pianov11 self config prevvocabsize config vocabsize config tiewordembeddings false model pop2pianoforconditionalgenerationconfigconfig totorchdevice eval model resizetokenembeddingsprevvocabsize 10 self parent assertequalmodel getinputembeddings weight shape0 prevvocabsize 10 self parent assertequalmodel getoutputembeddings weight shape0 prevvocabsize 10 self parent assertequalmodel config vocabsize prevvocabsize 10 def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config inputids decoderinputids attentionmask decoderattentionmask lmlabels configandinputs inputsdict inputids inputids attentionmask attentionmask decoderinputids decoderinputids decoderattentionmask decoderattentionmask usecache false return config inputsdict requiretorch class pop2pianomodeltestmodeltestermixin generationtestermixin pipelinetestermixin unittest testcase allmodelclasses pop2pianoforconditionalgeneration if istorchavailable else allgenerativemodelclasses pipelinemodelmapping automaticspeechrecognition pop2pianoforconditionalgeneration if istorchavailable else allparallelizablemodelclasses fxcompatible false testpruning false testresizeembeddings true testmodelparallel false isencoderdecoder true def setupself self modeltester pop2pianomodeltesterself self configtester configtesterself configclasspop2pianoconfig dmodel37 def testconfigself self configtester runcommontests def testshiftrightself configandinputs self modeltester prepareconfigandinputs self modeltester checkpreparelmlabelsviashiftleftconfigandinputs def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testmodelv11self configandinputs self modeltester prepareconfigandinputs check that gated gelu feed forward and different word embeddings work config configandinputs0 config tiewordembeddings false config feedforwardproj gatedgelu self modeltester createandcheckmodelconfig configandinputs1 def testconfigandmodelsilugatedself configandinputs self modeltester prepareconfigandinputs config configandinputs0 config feedforwardproj gatedsilu self modeltester createandcheckmodelconfigandinputs def testwithlmheadself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckwithlmheadconfigandinputs def testdecodermodelpastself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelpastconfigandinputs def testdecodermodelpastwithattnmaskself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelattentionmaskpastconfigandinputs def testdecodermodelpastwith3dattnmaskself config inputids decoderinputids attentionmask decoderattentionmask lmlabels self modeltester prepareconfigandinputs attentionmask idstensor self modeltester batchsize self modeltester encoderseqlength self modeltester encoderseqlength vocabsize2 decoderattentionmask idstensor self modeltester batchsize self modeltester decoderseqlength self modeltester decoderseqlength vocabsize2 self modeltester createandcheckdecodermodelattentionmaskpast config inputids decoderinputids attentionmask decoderattentionmask lmlabels def testdecodermodelpastwithlargeinputsself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckdecodermodelpastlargeinputsconfigandinputs def testencoderdecodersharedweightsself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckencoderdecodersharedweightsconfigandinputs unittest skipiftorchdevice cpu cant do half precision def testmodelfp16forwardself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelfp16forwardconfigandinputs def testv11resizeembeddingsself config self modeltester prepareconfigandinputs0 self modeltester checkresizeembeddingspop2pianov11config slow def testmodelfrompretrainedself for modelname in pop2pianopretrainedmodelarchivelist 1 model pop2pianoforconditionalgeneration frompretrainedmodelname self assertisnotnonemodel requireonnx unittest skipif istorch180 reasontest has a segmentation fault on torch 1 8 0 def testexporttoonnxself configandinputs self modeltester prepareconfigandinputs model pop2pianoforconditionalgenerationconfigandinputs0 totorchdevice with tempfile temporarydirectory as tmpdirname torch onnx export model configandinputs1 configandinputs3 configandinputs2 ftmpdirnamepop2pianotest onnx exportparamstrue opsetversion9 inputnamesinputids decoderinputids def testpasswithinputfeaturesself inputfeatures batchfeature inputfeatures torch rand75 100 512 typetorch float32 beatsteps torch randintsize1 955 low0 high100 typetorch float32 extrapolatedbeatstep torch randintsize1 900 low0 high100 typetorch float32 model pop2pianoforconditionalgeneration frompretrainedsweetcocoapop2piano modelopts model generateinputfeaturesinputfeaturesinputfeatures returndictingeneratetrue self assertequalmodelopts sequences ndim 2 def testpasswithbatchedinputfeaturesself inputfeatures batchfeature inputfeatures torch rand220 70 512 typetorch float32 beatsteps torch randintsize5 955 low0 high100 typetorch float32 extrapolatedbeatstep torch randintsize5 900 low0 high100 typetorch float32 attentionmask torch concatenate torch ones120 70 dtypetorch int32 torch zeros1 70 dtypetorch int32 torch ones50 70 dtypetorch int32 torch zeros1 70 dtypetorch int32 torch ones47 70 dtypetorch int32 torch zeros1 70 dtypetorch int32 axis0 attentionmaskbeatsteps torch ones5 955 typetorch int32 attentionmaskextrapolatedbeatstep torch ones5 900 typetorch int32 model pop2pianoforconditionalgeneration frompretrainedsweetcocoapop2piano modelopts model generate inputfeaturesinputfeaturesinputfeatures attentionmaskinputfeaturesattentionmask returndictingeneratetrue self assertequalmodelopts sequences ndim 2 requiretorch class pop2pianomodelintegrationtestsunittest testcase slow def testmelconditionerintegrationself composer composer1 model pop2pianoforconditionalgeneration frompretrainedsweetcocoapop2piano inputembeds torch ones10 100 512 composervalue model generationconfig composertofeaturetokencomposer composervalue torch tensorcomposervalue composervalue composervalue repeatinputembeds size0 outputs model melconditioner inputembeds composervalue minmodel generationconfig composertofeaturetoken values check shape self assertequaloutputs size torch size10 101 512 check values expectedoutputs torch tensor 1 0475305318832397 0 29052114486694336 0 47778210043907166 1 0 1 0 1 0 1 0 1 0 1 0 self asserttruetorch allcloseoutputs0 3 3 expectedoutputs atol1e4 slow requireessentia requirelibrosa requirescipy def testfullmodelintegrationself if islibrosaavailable and isscipyavailable and isessentiaavailable and istorchavailable from transformers import pop2pianoprocessor speechinput1 np zeros1000000 dtypenp float32 samplingrate 44100 processor pop2pianoprocessor frompretrainedsweetcocoapop2piano inputfeatures processor featureextractor speechinput1 samplingratesamplingrate returntensorspt model pop2pianoforconditionalgeneration frompretrainedsweetcocoapop2piano outputs model generate inputfeaturesinputfeaturesinputfeatures returndictingeneratetrue sequences check for shapes self assertequaloutputs size0 70 check for values self assertequaloutputs0 2 detach cpu numpy tolist 0 1 this is the test for a real music from kpop genre slow requireessentia requirelibrosa requirescipy def testrealmusicself if islibrosaavailable and isscipyavailable and isessentiaavailable and istorchavailable from transformers import pop2pianofeatureextractor pop2pianotokenizer model pop2pianoforconditionalgeneration frompretrainedsweetcocoapop2piano model eval featureextractor pop2pianofeatureextractor frompretrainedsweetcocoapop2piano tokenizer pop2pianotokenizer frompretrainedsweetcocoapop2piano ds loaddatasetsweetcocoapop2pianoci splittest outputfe featureextractor dsaudio0array samplingratedsaudio0samplingrate returntensorspt outputmodel model generateinputfeaturesoutputfeinputfeatures composercomposer1 outputtokenizer tokenizer batchdecodetokenidsoutputmodel featureextractoroutputoutputfe prettymidiobject outputtokenizerprettymidiobjects0 checking if no of notes are same self assertequallenprettymidiobject instruments0 notes 59 predictedtimings for i in prettymidiobject instruments0 notes predictedtimings appendi start checking note start timingsfirst 6 expectedstarttimings 0 4876190423965454 0 7314285635948181 0 9752380847930908 1 4396371841430664 1 6718367338180542 1 904036283493042 np allcloseexpectedstarttimings predictedtimings 6 checking note end timingslast 6 expectedendtimings 12 341403007507324 12 567797183990479 12 567797183990479 12 567797183990479 12 794191360473633 12 794191360473633 np allcloseexpectedendtimings predictedtimings6 coding utf 8 2023 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 pop2piano model for common tests for common tests pop2piano forces 100 extra tokens make sure that lm_labels are correctly padded from the right add causal pad token mask first item items before diagonal pad items after diagonal all items after square there should be num_layers key value embeddings stored in decoder_past there should be a self attn key a self attn value a cross attn key and a cross attn value stored in each decoder_past tuple first forward pass create hypothetical next token and extent to next_input_ids append to next input_ids and select random slice test that outputs are equal for slice create attention mask first forward pass create hypothetical next token and extent to next_input_ids change a random masked slice from input_ids append to next input_ids and attn_mask get two different outputs select random slice test that outputs are equal for slice first forward pass create hypothetical multiple next token and extent to next_input_ids append to next input_ids and select random slice test that outputs are equal for slice load state dict copies weights but does not tie them check that models has less parameters check that outputs are equal check that outputs after saving and loading are equal check that models has less parameters check that outputs are equal check that gated gelu feed forward and different word embeddings work check shape check values check for shapes check for values this is the test for a real music from k pop genre checking if no of notes are same checking note start timings first 6 checking note end timings last 6
import copy
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import Pop2PianoConfig
from transformers.feature_extraction_utils import BatchFeature
from transformers.testing_utils import (
    require_essentia,
    require_librosa,
    require_onnx,
    require_scipy,
    require_torch,
    slow,
    torch_device,
)
from transformers.utils import is_essentia_available, is_librosa_available, is_scipy_available, is_torch_available

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


if is_torch_available():
    import torch

    from transformers import Pop2PianoForConditionalGeneration
    from transformers.models.pop2piano.modeling_pop2piano import POP2PIANO_PRETRAINED_MODEL_ARCHIVE_LIST
    from transformers.pytorch_utils import is_torch_1_8_0

else:
    is_torch_1_8_0 = False


@require_torch
class Pop2PianoModelTester:
    def __init__(
        self,
        parent,
        vocab_size=99,
        batch_size=13,
        encoder_seq_length=7,
        decoder_seq_length=9,
        
        is_training=False,
        use_attention_mask=True,
        use_labels=True,
        hidden_size=64,
        num_hidden_layers=5,
        num_attention_heads=4,
        d_ff=37,
        relative_attention_num_buckets=8,
        dropout_rate=0.1,
        initializer_factor=0.002,
        eos_token_id=1,
        pad_token_id=0,
        decoder_start_token_id=0,
        scope=None,
        decoder_layers=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.encoder_seq_length = encoder_seq_length
        self.decoder_seq_length = decoder_seq_length
        
        self.seq_length = self.decoder_seq_length
        self.is_training = is_training
        self.use_attention_mask = use_attention_mask
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.d_ff = d_ff
        self.relative_attention_num_buckets = relative_attention_num_buckets
        self.dropout_rate = dropout_rate
        self.initializer_factor = initializer_factor
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.decoder_start_token_id = decoder_start_token_id
        self.scope = None
        self.decoder_layers = decoder_layers

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.encoder_seq_length], self.vocab_size)
        decoder_input_ids = ids_tensor([self.batch_size, self.decoder_seq_length], self.vocab_size)

        attention_mask = None
        decoder_attention_mask = None
        if self.use_attention_mask:
            attention_mask = ids_tensor([self.batch_size, self.encoder_seq_length], vocab_size=2)
            decoder_attention_mask = ids_tensor([self.batch_size, self.decoder_seq_length], vocab_size=2)

        lm_labels = (
            ids_tensor([self.batch_size, self.decoder_seq_length], self.vocab_size) if self.use_labels else None
        )

        return self.get_config(), input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels

    def get_pipeline_config(self):
        return Pop2PianoConfig(
            vocab_size=166,  
            d_model=self.hidden_size,
            d_ff=self.d_ff,
            d_kv=self.hidden_size // self.num_attention_heads,
            num_layers=self.num_hidden_layers,
            num_decoder_layers=self.decoder_layers,
            num_heads=self.num_attention_heads,
            relative_attention_num_buckets=self.relative_attention_num_buckets,
            dropout_rate=self.dropout_rate,
            initializer_factor=self.initializer_factor,
            eos_token_id=self.eos_token_id,
            bos_token_id=self.pad_token_id,
            pad_token_id=self.pad_token_id,
            decoder_start_token_id=self.decoder_start_token_id,
        )

    def get_config(self):
        return Pop2PianoConfig(
            vocab_size=self.vocab_size,
            d_model=self.hidden_size,
            d_ff=self.d_ff,
            d_kv=self.hidden_size // self.num_attention_heads,
            num_layers=self.num_hidden_layers,
            num_decoder_layers=self.decoder_layers,
            num_heads=self.num_attention_heads,
            relative_attention_num_buckets=self.relative_attention_num_buckets,
            dropout_rate=self.dropout_rate,
            initializer_factor=self.initializer_factor,
            eos_token_id=self.eos_token_id,
            bos_token_id=self.pad_token_id,
            pad_token_id=self.pad_token_id,
            decoder_start_token_id=self.decoder_start_token_id,
        )

    def check_prepare_lm_labels_via_shift_left(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config)
        model.to(torch_device)
        model.eval()

        
        lm_labels.masked_fill_((lm_labels == self.decoder_start_token_id), self.eos_token_id)

        
        triangular_mask = torch.tril(lm_labels.new_ones(lm_labels.shape)).logical_not()
        lm_labels.masked_fill_(triangular_mask, self.pad_token_id)
        decoder_input_ids = model._shift_right(lm_labels)

        for i, (decoder_input_ids_slice, lm_labels_slice) in enumerate(zip(decoder_input_ids, lm_labels)):
            
            self.parent.assertEqual(decoder_input_ids_slice[0].item(), self.decoder_start_token_id)
            if i < decoder_input_ids_slice.shape[-1]:
                if i < decoder_input_ids.shape[-1] - 1:
                    
                    self.parent.assertListEqual(
                        decoder_input_ids_slice[1 : i + 1].tolist(), lm_labels_slice[:i].tolist()
                    )
                
                if i < decoder_input_ids.shape[-1] - 2:
                    self.parent.assertListEqual(
                        decoder_input_ids_slice[i + 2 :].tolist(), lm_labels_slice[i + 1 : -1].tolist()
                    )
            else:
                
                self.parent.assertListEqual(decoder_input_ids_slice[1:].tolist(), lm_labels_slice[:-1].tolist())

    def create_and_check_model(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config)
        model.to(torch_device)
        model.eval()
        result = model(
            input_ids=input_ids,
            decoder_input_ids=decoder_input_ids,
            attention_mask=attention_mask,
            decoder_attention_mask=decoder_attention_mask,
        )
        result = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)
        decoder_past = result.past_key_values
        encoder_output = result.encoder_last_hidden_state

        self.parent.assertEqual(encoder_output.size(), (self.batch_size, self.encoder_seq_length, self.hidden_size))
        
        self.parent.assertEqual(len(decoder_past), config.num_layers)
        
        self.parent.assertEqual(len(decoder_past[0]), 4)

    def create_and_check_with_lm_head(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).to(torch_device).eval()
        outputs = model(
            input_ids=input_ids,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            labels=lm_labels,
        )
        self.parent.assertEqual(len(outputs), 4)
        self.parent.assertEqual(outputs["logits"].size(), (self.batch_size, self.decoder_seq_length, self.vocab_size))
        self.parent.assertEqual(outputs["loss"].size(), ())

    def create_and_check_decoder_model_past(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).get_decoder().to(torch_device).eval()
        
        outputs = model(input_ids, use_cache=True)
        outputs_use_cache_conf = model(input_ids)
        outputs_no_past = model(input_ids, use_cache=False)

        self.parent.assertTrue(len(outputs) == len(outputs_use_cache_conf))
        self.parent.assertTrue(len(outputs) == len(outputs_no_past) + 1)

        output, past_key_values = outputs.to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)

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

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

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

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

    def create_and_check_decoder_model_attention_mask_past(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).get_decoder()
        model.to(torch_device)
        model.eval()

        
        attn_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)

        half_seq_length = input_ids.shape[-1] // 2
        attn_mask[:, half_seq_length:] = 0

        
        output, past_key_values = model(input_ids, attention_mask=attn_mask, use_cache=True).to_tuple()

        
        next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size)

        
        random_seq_idx_to_change = ids_tensor((1,), half_seq_length).item() + 1
        random_other_next_tokens = ids_tensor((self.batch_size, 1), config.vocab_size).squeeze(-1)
        input_ids[:, -random_seq_idx_to_change] = random_other_next_tokens

        
        next_input_ids = torch.cat([input_ids, next_tokens], dim=-1)
        attn_mask = torch.cat(
            [attn_mask, torch.ones((attn_mask.shape[0], 1), dtype=torch.long, device=torch_device)],
            dim=1,
        )

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

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

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

    def create_and_check_decoder_model_past_large_inputs(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).get_decoder().to(torch_device).eval()
        
        outputs = model(input_ids, attention_mask=attention_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

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

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

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

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

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

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

    def create_and_check_generate_with_past_key_values(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).to(torch_device).eval()
        torch.manual_seed(0)
        output_without_past_cache = model.generate(
            input_ids[:1], num_beams=2, max_length=5, do_sample=True, use_cache=False
        )
        torch.manual_seed(0)
        output_with_past_cache = model.generate(input_ids[:1], num_beams=2, max_length=5, do_sample=True)
        self.parent.assertTrue(torch.all(output_with_past_cache == output_without_past_cache))

    def create_and_check_model_fp16_forward(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        model = Pop2PianoForConditionalGeneration(config=config).to(torch_device).half().eval()
        output = model(input_ids, decoder_input_ids=input_ids, attention_mask=attention_mask)[
            "encoder_last_hidden_state"
        ]
        self.parent.assertFalse(torch.isnan(output).any().item())

    def create_and_check_encoder_decoder_shared_weights(
        self,
        config,
        input_ids,
        decoder_input_ids,
        attention_mask,
        decoder_attention_mask,
        lm_labels,
    ):
        for model_class in [Pop2PianoForConditionalGeneration]:
            torch.manual_seed(0)
            model = model_class(config=config).to(torch_device).eval()
            
            model.encoder.load_state_dict(model.decoder.state_dict(), strict=False)

            torch.manual_seed(0)
            tied_config = copy.deepcopy(config)
            tied_config.tie_encoder_decoder = True
            tied_model = model_class(config=tied_config).to(torch_device).eval()

            model_result = model(
                input_ids=input_ids,
                decoder_input_ids=decoder_input_ids,
                attention_mask=attention_mask,
                decoder_attention_mask=decoder_attention_mask,
            )

            tied_model_result = tied_model(
                input_ids=input_ids,
                decoder_input_ids=decoder_input_ids,
                attention_mask=attention_mask,
                decoder_attention_mask=decoder_attention_mask,
            )

            
            self.parent.assertLess(
                sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters())
            )
            random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()

            
            self.parent.assertTrue(
                torch.allclose(
                    model_result[0][0, :, random_slice_idx], tied_model_result[0][0, :, random_slice_idx], atol=1e-4
                )
            )

            
            with tempfile.TemporaryDirectory() as tmpdirname:
                tied_model.save_pretrained(tmpdirname)
                tied_model = model_class.from_pretrained(tmpdirname)
                tied_model.to(torch_device)
                tied_model.eval()

                
                self.parent.assertLess(
                    sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters())
                )
                random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()

                tied_model_result = tied_model(
                    input_ids=input_ids,
                    decoder_input_ids=decoder_input_ids,
                    attention_mask=attention_mask,
                    decoder_attention_mask=decoder_attention_mask,
                )

                
                self.parent.assertTrue(
                    torch.allclose(
                        model_result[0][0, :, random_slice_idx],
                        tied_model_result[0][0, :, random_slice_idx],
                        atol=1e-4,
                    )
                )

    def check_resize_embeddings_pop2piano_v1_1(
        self,
        config,
    ):
        prev_vocab_size = config.vocab_size

        config.tie_word_embeddings = False
        model = Pop2PianoForConditionalGeneration(config=config).to(torch_device).eval()
        model.resize_token_embeddings(prev_vocab_size - 10)

        self.parent.assertEqual(model.get_input_embeddings().weight.shape[0], prev_vocab_size - 10)
        self.parent.assertEqual(model.get_output_embeddings().weight.shape[0], prev_vocab_size - 10)
        self.parent.assertEqual(model.config.vocab_size, prev_vocab_size - 10)

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (
            config,
            input_ids,
            decoder_input_ids,
            attention_mask,
            decoder_attention_mask,
            lm_labels,
        ) = config_and_inputs

        inputs_dict = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "decoder_input_ids": decoder_input_ids,
            "decoder_attention_mask": decoder_attention_mask,
            "use_cache": False,
        }
        return config, inputs_dict


@require_torch
class Pop2PianoModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (Pop2PianoForConditionalGeneration,) if is_torch_available() else ()
    all_generative_model_classes = ()
    pipeline_model_mapping = (
        {"automatic-speech-recognition": Pop2PianoForConditionalGeneration} if is_torch_available() else {}
    )
    all_parallelizable_model_classes = ()
    fx_compatible = False
    test_pruning = False
    test_resize_embeddings = True
    test_model_parallel = False
    is_encoder_decoder = True

    def setUp(self):
        self.model_tester = Pop2PianoModelTester(self)
        self.config_tester = ConfigTester(self, config_class=Pop2PianoConfig, d_model=37)

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

    def test_shift_right(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_prepare_lm_labels_via_shift_left(*config_and_inputs)

    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_model_v1_1(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        
        config = config_and_inputs[0]
        config.tie_word_embeddings = False
        config.feed_forward_proj = "gated-gelu"
        self.model_tester.create_and_check_model(config, *config_and_inputs[1:])

    def test_config_and_model_silu_gated(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        config = config_and_inputs[0]
        config.feed_forward_proj = "gated-silu"
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_with_lm_head(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_with_lm_head(*config_and_inputs)

    def test_decoder_model_past(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_past(*config_and_inputs)

    def test_decoder_model_past_with_attn_mask(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_decoder_model_attention_mask_past(*config_and_inputs)

    def test_decoder_model_past_with_3d_attn_mask(self):
        (
            config,
            input_ids,
            decoder_input_ids,
            attention_mask,
            decoder_attention_mask,
            lm_labels,
        ) = self.model_tester.prepare_config_and_inputs()

        attention_mask = ids_tensor(
            [self.model_tester.batch_size, self.model_tester.encoder_seq_length, self.model_tester.encoder_seq_length],
            vocab_size=2,
        )
        decoder_attention_mask = ids_tensor(
            [self.model_tester.batch_size, self.model_tester.decoder_seq_length, self.model_tester.decoder_seq_length],
            vocab_size=2,
        )

        self.model_tester.create_and_check_decoder_model_attention_mask_past(
            config,
            input_ids,
            decoder_input_ids,
            attention_mask,
            decoder_attention_mask,
            lm_labels,
        )

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

    def test_encoder_decoder_shared_weights(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_encoder_decoder_shared_weights(*config_and_inputs)

    @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
    def test_model_fp16_forward(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model_fp16_forward(*config_and_inputs)

    def test_v1_1_resize_embeddings(self):
        config = self.model_tester.prepare_config_and_inputs()[0]
        self.model_tester.check_resize_embeddings_pop2piano_v1_1(config)

    @slow
    def test_model_from_pretrained(self):
        for model_name in POP2PIANO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = Pop2PianoForConditionalGeneration.from_pretrained(model_name)
            self.assertIsNotNone(model)

    @require_onnx
    @unittest.skipIf(
        is_torch_1_8_0,
        reason="Test has a segmentation fault on torch 1.8.0",
    )
    def test_export_to_onnx(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        model = Pop2PianoForConditionalGeneration(config_and_inputs[0]).to(torch_device)
        with tempfile.TemporaryDirectory() as tmpdirname:
            torch.onnx.export(
                model,
                (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2]),
                f"{tmpdirname}/Pop2Piano_test.onnx",
                export_params=True,
                opset_version=9,
                input_names=["input_ids", "decoder_input_ids"],
            )

    def test_pass_with_input_features(self):
        input_features = BatchFeature(
            {
                "input_features": torch.rand((75, 100, 512)).type(torch.float32),
                "beatsteps": torch.randint(size=(1, 955), low=0, high=100).type(torch.float32),
                "extrapolated_beatstep": torch.randint(size=(1, 900), low=0, high=100).type(torch.float32),
            }
        )
        model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano")
        model_opts = model.generate(input_features=input_features["input_features"], return_dict_in_generate=True)

        self.assertEqual(model_opts.sequences.ndim, 2)

    def test_pass_with_batched_input_features(self):
        input_features = BatchFeature(
            {
                "input_features": torch.rand((220, 70, 512)).type(torch.float32),
                "beatsteps": torch.randint(size=(5, 955), low=0, high=100).type(torch.float32),
                "extrapolated_beatstep": torch.randint(size=(5, 900), low=0, high=100).type(torch.float32),
                "attention_mask": torch.concatenate(
                    [
                        torch.ones([120, 70], dtype=torch.int32),
                        torch.zeros([1, 70], dtype=torch.int32),
                        torch.ones([50, 70], dtype=torch.int32),
                        torch.zeros([1, 70], dtype=torch.int32),
                        torch.ones([47, 70], dtype=torch.int32),
                        torch.zeros([1, 70], dtype=torch.int32),
                    ],
                    axis=0,
                ),
                "attention_mask_beatsteps": torch.ones((5, 955)).type(torch.int32),
                "attention_mask_extrapolated_beatstep": torch.ones((5, 900)).type(torch.int32),
            }
        )
        model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano")
        model_opts = model.generate(
            input_features=input_features["input_features"],
            attention_mask=input_features["attention_mask"],
            return_dict_in_generate=True,
        )

        self.assertEqual(model_opts.sequences.ndim, 2)


@require_torch
class Pop2PianoModelIntegrationTests(unittest.TestCase):
    @slow
    def test_mel_conditioner_integration(self):
        composer = "composer1"
        model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano")
        input_embeds = torch.ones([10, 100, 512])

        composer_value = model.generation_config.composer_to_feature_token[composer]
        composer_value = torch.tensor(composer_value)
        composer_value = composer_value.repeat(input_embeds.size(0))
        outputs = model.mel_conditioner(
            input_embeds, composer_value, min(model.generation_config.composer_to_feature_token.values())
        )

        
        self.assertEqual(outputs.size(), torch.Size([10, 101, 512]))

        
        EXPECTED_OUTPUTS = torch.tensor(
            [[1.0475305318832397, 0.29052114486694336, -0.47778210043907166], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]
        )

        self.assertTrue(torch.allclose(outputs[0, :3, :3], EXPECTED_OUTPUTS, atol=1e-4))

    @slow
    @require_essentia
    @require_librosa
    @require_scipy
    def test_full_model_integration(self):
        if is_librosa_available() and is_scipy_available() and is_essentia_available() and is_torch_available():
            from transformers import Pop2PianoProcessor

            speech_input1 = np.zeros([1_000_000], dtype=np.float32)
            sampling_rate = 44_100

            processor = Pop2PianoProcessor.from_pretrained("sweetcocoa/pop2piano")
            input_features = processor.feature_extractor(
                speech_input1, sampling_rate=sampling_rate, return_tensors="pt"
            )

            model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano")
            outputs = model.generate(
                input_features=input_features["input_features"], return_dict_in_generate=True
            ).sequences

            
            self.assertEqual(outputs.size(0), 70)

            
            self.assertEqual(outputs[0, :2].detach().cpu().numpy().tolist(), [0, 1])

    
    @slow
    @require_essentia
    @require_librosa
    @require_scipy
    def test_real_music(self):
        if is_librosa_available() and is_scipy_available() and is_essentia_available() and is_torch_available():
            from transformers import Pop2PianoFeatureExtractor, Pop2PianoTokenizer

            model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano")
            model.eval()
            feature_extractor = Pop2PianoFeatureExtractor.from_pretrained("sweetcocoa/pop2piano")
            tokenizer = Pop2PianoTokenizer.from_pretrained("sweetcocoa/pop2piano")
            ds = load_dataset("sweetcocoa/pop2piano_ci", split="test")

            output_fe = feature_extractor(
                ds["audio"][0]["array"], sampling_rate=ds["audio"][0]["sampling_rate"], return_tensors="pt"
            )
            output_model = model.generate(input_features=output_fe["input_features"], composer="composer1")
            output_tokenizer = tokenizer.batch_decode(token_ids=output_model, feature_extractor_output=output_fe)
            pretty_midi_object = output_tokenizer["pretty_midi_objects"][0]

            
            self.assertEqual(len(pretty_midi_object.instruments[0].notes), 59)
            predicted_timings = []
            for i in pretty_midi_object.instruments[0].notes:
                predicted_timings.append(i.start)

            
            EXPECTED_START_TIMINGS = [
                0.4876190423965454,
                0.7314285635948181,
                0.9752380847930908,
                1.4396371841430664,
                1.6718367338180542,
                1.904036283493042,
            ]

            np.allclose(EXPECTED_START_TIMINGS, predicted_timings[:6])

            
            EXPECTED_END_TIMINGS = [
                12.341403007507324,
                12.567797183990479,
                12.567797183990479,
                12.567797183990479,
                12.794191360473633,
                12.794191360473633,
            ]

            np.allclose(EXPECTED_END_TIMINGS, predicted_timings[-6:])