o Zh@s`dZddlZddlZddlmZmZmZddlZddlmZddl m Z ddl m Z ddl mZdd lmZmZmZdd lmZdd lmZdd lmZmZmZdd lmZddlmZmZm Z ddl!m"Z"m#Z#m$Z$m%Z%m&Z&ddl'm(Z(e#rddl)m*Z*ddl+m,Z,e&-e.Z/dZ0zddl1m2Z2dZ0e/3dWne4yYne5ye/6dYnwGdddej7Z8e0se2Z8e9e8Gdddej7Z:Gdddej7Z;Gdd d ej7ZGd%d&d&ej7Z?Gd'd(d(ej7Z@e"Gd)d*d*eZAGd+d,d,eAZBGd-d.d.ej7ZCe"d/d0Gd1d2d2eAeZDd2d*gZEdS)3zPyTorch Pop2Piano model.N)OptionalTupleUnion)nn)CrossEntropyLoss)GenerationConfig)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)AttentionMaskConverter)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsSeq2SeqLMOutput)PreTrainedModel)ALL_LAYERNORM_LAYERS find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringis_torch_flex_attn_availableis_torch_fx_proxyis_torchdynamo_compilinglogging)Pop2PianoConfig) BlockMask)make_flex_block_causal_maskT) FusedRMSNormFzVDiscovered apex.normalization.FusedRMSNorm - will use it instead of Pop2PianoLayerNormzIDiscovered apex but it failed to load, falling back to Pop2PianoLayerNormcs&eZdZdfdd ZddZZS)Pop2PianoLayerNormư>cs&ttt||_||_dS)zj Construct a layernorm module in the Pop2Piano style. No bias and no subtraction of mean. N)super__init__r Parametertorchonesweightvariance_epsilon)selfZ hidden_sizeeps __class___/var/www/auris/lib/python3.10/site-packages/transformers/models/pop2piano/modeling_pop2piano.pyr#@s  zPop2PianoLayerNorm.__init__cCs\|tjdjddd}|t||j}|jjtj tj fvr)||jj}|j|S)NT)Zkeepdim) tor%Zfloat32powmeanZrsqrtr(r'dtypefloat16Zbfloat16)r) hidden_statesZvariancer-r-r.forwardHs  zPop2PianoLayerNorm.forward)r!)__name__ __module__ __qualname__r#r7 __classcell__r-r-r+r.r ?sr c*eZdZdeffdd ZddZZS)Pop2PianoDenseActDenseconfigcsTttj|j|jdd|_tj|j|jdd|_t|j |_ t |j |_ dSNFbias)r"r#rLineard_modeld_ffwiwoDropout dropout_ratedropoutr dense_act_fnactr)r>r+r-r.r#`s zPop2PianoDenseActDense.__init__cCsl||}||}||}t|jjtjr/|j|jjjkr/|jjjtj kr/| |jjj}||}|SN) rErKrI isinstancerFr'r%Tensorr4int8r1)r)r6r-r-r.r7gs    zPop2PianoDenseActDense.forwardr8r9r:rr#r7r;r-r-r+r.r=_sr=cr<)Pop2PianoDenseGatedActDenser>csjttj|j|jdd|_tj|j|jdd|_tj|j|jdd|_t |j |_ t |j |_dSr?)r"r#rrBrCrDwi_0wi_1rFrGrHrIr rJrKrLr+r-r.r#ws z$Pop2PianoDenseGatedActDense.__init__cCsz|||}||}||}||}t|jjtjr6|j |jjj kr6|jjj tj kr6| |jjj }||}|SrM) rKrSrTrIrNrFr'r%rOr4rPr1)r)r6Z hidden_geluZ hidden_linearr-r-r.r7s   z#Pop2PianoDenseGatedActDense.forwardrQr-r-r+r.rRvsrRcr<)Pop2PianoLayerFFr>csJt|jrt||_nt||_t|j|jd|_ t |j |_ dS)Nr*)r"r#Z is_gated_actrRDenseReluDenser=r rClayer_norm_epsilon layer_normrrGrHrIrLr+r-r.r#s   zPop2PianoLayerFF.__init__cCs&||}||}|||}|SrM)rYrWrI)r)r6Zforwarded_statesr-r-r.r7s  zPop2PianoLayerFF.forwardrQr-r-r+r.rUs rUcsleZdZ  ddedeeffdd ZddZedd d Z dddZ         dddZ Z S)Pop2PianoAttentionFNr> layer_idxcst|j|_||_|j|_|j|_|j|_|j|_|j |_ |j |_ |j |j|_ ||_|dur@|jr@td|jjdtj|j|j dd|_tj|j|j dd|_tj|j|j dd|_tj|j |jdd|_|jrxt|j|j |_t|_d|_dS)NzInstantiating a decoder z without passing `layer_idx` is not recommended and will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.Fr@)r"r# is_decoderhas_relative_attention_biasrelative_attention_num_bucketsrelative_attention_max_distancerCd_kvkey_value_proj_dim num_headsn_headsrHrI inner_dimr[logger warning_oncer,r8rrBqkvo Embeddingrelative_attention_biasset pruned_headsgradient_checkpointingr)r>r]r[r+r-r.r#s.  zPop2PianoAttention.__init__cCst|dkrdSt||j|j|j\}}t|j||_t|j||_t|j||_t|j |dd|_ |jt||_|j|j|_ |j ||_dS)Nrrdim) lenrrcrarnrrgrhrirjrdunion)r)Zheadsindexr-r-r. prune_headss zPop2PianoAttention.prune_headsT cCsd}|r|d}||dktj|7}t|}n t|t| }|d}||k}|t||t||||tj}t|t ||d}|t |||7}|S)a Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 Translate relative position to a bucket number for relative attention. The relative position is defined as memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should allow for more graceful generalization to longer sequences than the model has been trained on Args: relative_position: an int32 Tensor bidirectional: a boolean - whether the attention is bidirectional num_buckets: an integer max_distance: an integer Returns: a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) rr/r) r1r%longabsminZ zeros_likelogfloatmathZ full_likewhere)relative_position bidirectional num_buckets max_distanceZrelative_bucketsZ max_exactZis_smallZrelative_position_if_larger-r-r._relative_position_buckets*  z,Pop2PianoAttention._relative_position_bucketc Cs|dur |jjj}|durtj|tj|ddddf}n |dddf|}tj|tj|ddddf}||}|j||j |j |j d}||} | gd d} | S)z%Compute binned relative position biasN)r4device)rrr)r/rrr) rlr'rr%arangeryr1rr\r^r_Zpermute unsqueeze) r) query_length key_lengthrcache_positionZcontext_positionZmemory_positionrZrelative_position_bucketvaluesr-r-r. compute_bias s    zPop2PianoAttention.compute_biasc Cs|jdd\} } |du} ||}|| d|j|jdd}|dur4|j|j}| r1|j }n|j }| r8|n|}| rO|durO|rO|j |j}|j |j}nE| |}||}|| d|j|jdd}|| d|j|jdd}|dur| s}| nd} ||||jd| i\}}| rd|j|j<t||dd}|dur|jd}|dur|n| dd}|jstjd|j| |f|j|jd }|jr|jrd|_n|j|||j| d }|dddd| dddf}|dur|ddddddd|jdf}||}|jr%t|jd}d |t|j<|dd|f}n|}||7}tjj |!dd "|}tjj#||j#|jd }|durL||}t||}|dd$}|| d|j%}|&|}|||f}| rt||f}|S)z Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). Nr/r0rrTr)rr4)rrrrq)ptraining)'shapergviewrcraZ transpose is_updatedgetr[Zcross_attention_cacheself_attention_cacheZ key_cacheZ value_cacherhriupdater%matmulr]Zzerosrr4rorZ requires_gradrrnr&listboolrZ functionalZsoftmaxr}Ztype_asrI contiguousrdrj)r)r6maskkey_value_states position_biaspast_key_valuelayer_head_maskr use_cacheoutput_attentionsr batch_size seq_lengthZis_cross_attentionZ query_statesrZcurr_past_key_valueZcurrent_statesZ key_statesZ value_statesZscoresrZreal_seq_length causal_maskZposition_bias_maskedZ attn_weightsZ attn_outputoutputsr-r-r.r7sx           " &     zPop2PianoAttention.forwardFN)Trwrx)NN) NNNNNNFFN) r8r9r:rrintr#rv staticmethodrrr7r;r-r-r+r.rZs,#  /rZcs@eZdZddeeffdd Z       d ddZZS) Pop2PianoLayerSelfAttentionFNr[cs>tt|||d|_t|j|jd|_t |j |_ dS)Nr]r[rV) r"r#rZ SelfAttentionr rCrXrYrrGrHrIrpr+r-r.r#s z$Pop2PianoLayerSelfAttention.__init__c CsL||} |j| |||||||d} ||| d}|f| dd} | S)N)rrrrrrrrr)rYrrI) r)r6attention_maskrrrrrrnormed_hidden_statesattention_outputrr-r-r.r7s  z#Pop2PianoLayerSelfAttention.forwardr)NNNNFFNr8r9r:rrr#r7r;r-r-r+r.rs rcsBeZdZddeeffdd Z        d ddZZS) Pop2PianoLayerCrossAttentionNr[cs>tt|d|d|_t|j|jd|_t |j |_ dS)NFrrV) r"r#rZEncDecAttentionr rCrXrYrrGrHrI)r)r>r[r+r-r.r#s z%Pop2PianoLayerCrossAttention.__init__Fc  CsP||} |j| |||||||| | d } ||| d} | f| dd}|S)N) rrrrrrrrrrr)rYrrI)r)r6rrrrrrrrrrrZ layer_outputrr-r-r.r7s  z$Pop2PianoLayerCrossAttention.forwardrM)NNNNFNFNrr-r-r+r.rs rcsJeZdZd deeffdd Z            d ddZZS) Pop2PianoBlockFNr[cs`t|j|_t|_|jt|||d|jr&|jt||d|jt |dS)Nr)r[) r"r#r\r ModuleListlayerappendrrrUrpr+r-r.r#s  zPop2PianoBlock.__init__Tc Cs|jd||||| | | | d}|dd\}} |dd}|jtjkrDtt|t|jjdt|jj}tj || |d}|j oJ|du}|r|jd|||||| | dd| | d }|dd\}} |jtjkrtt|t|jjdt|jj}tj || |d}||dd}|jd|}|jtjkrtt|t|jjdt|jj}tj || |d}|f}| r|| f|}|S||}|S) Nr)rrrrrrrr/i)r{maxrr0)rrrrrrrr) rr4r%r5risinfanyfinforclampr\)r)r6rrencoder_hidden_statesencoder_attention_maskencoder_decoder_position_biasrcross_attn_layer_head_maskrrr return_dictrZself_attention_outputsZattention_outputsZ clamp_valueZdo_cross_attentionZcross_attention_outputsrr-r-r.r7sn           zPop2PianoBlock.forwardr) NNNNNNNNFFTNrr-r-r+r.rsrc@s@eZdZeZdZdZdZdZdZ dgZ dgZ ddZ dd Z d S) Pop2PianoPreTrainedModelZ transformerFTrrFcCs|jj}t|tr|jj|ddSt|tr'|jjjj d|dddSt|t rS|j jjj d|ddt |drO|jj sQ|jjjj d|dddSdSdSt|tr|jjjj d||jjddt |jdr{|jjdur{|jjj|jjjj d||jjddt |jdr|jjdur|jjjdSdSdSt|tr|jjjj d||jjddt |jdr|jjdur|jjj|jjjj d||jjddt |jdr|jjdur|jjj|jjjj d||jjddt |jdr|jjdur|jjjdSdSdSt|tr}|jj}|jj}|jj}|jjjj d|||dd|jjjj d||dd|jjjj d||dd|jjjj d|||dd|j r|j!jjj d||dddSdSdS)zInitialize the weights?)r3Zstdlm_headrAN)"r>Zinitializer_factorrNr r'dataZfill_Pop2PianoConcatEmbeddingToMel embeddingZnormal_!Pop2PianoForConditionalGenerationsharedhasattrtie_word_embeddingsrr=rErCrAZzero_rFrDrRrSrTrZr`rbrgrhrirjr]rl)r)modulefactorrCrarcr-r-r. _init_weightsIsR               z&Pop2PianoPreTrainedModel._init_weightscCs|jj}|jj}|durtdt|r1t|jddd|}tj||dddfgdd}n| |j}|dddf |dddf<||d<|durStd| |d k||S) Nzoself.model.config.decoder_start_token_id has to be defined. In Pop2Piano it is usually set to the pad_token_id.r0)r.rqr).rz1self.model.config.pad_token_id has to be defined.) r>decoder_start_token_id pad_token_id ValueErrorrr%fullrcatZ new_zeroscloneZ masked_fill_)r) input_idsrrZshifted_input_idsr-r-r. _shift_rightws    z%Pop2PianoPreTrainedModel._shift_rightN)r8r9r:rZ config_classZbase_model_prefixZis_parallelizableZsupports_gradient_checkpointingZ_supports_cache_classZ_supports_static_cacheZ_no_split_modulesZ_keep_in_fp32_modulesrrr-r-r-r.r>s .rc seZdZdfdd ZddZddZ             ddd Z dd eej d fd ej dej de de f ddZ e d ej dededejdej def ddZZS)Pop2PianoStackNcsxt||_j|_tfddtjD|_t j j d|_ t j|_|d|_d|_d|_dS)Ncs"g|] }tt|dk|dqS)rr)rr).0ir>r-r. sz+Pop2PianoStack.__init__..rVF)r"r# embed_tokensr\rrrange num_layersblockr rCrXfinal_layer_normrGrHrI post_initZmodel_parallelZ device_mapro)r)r>rr+rr.r#s   zPop2PianoStack.__init__cC|jSrMrr)r-r-r.get_input_embeddingsz#Pop2PianoStack.get_input_embeddingscC ||_dSrMrr)Znew_embeddingsr-r-r.set_input_embeddings z#Pop2PianoStack.set_input_embeddingsc)Cs| dur| n|jj} | dur| n|jj} | dur| n|jj} | dur$| n|jj} |durB|durB|jr5dnd}td|d|d|durS|}|d|d}n|dur`|dd}n|jrednd}td|d|d |j r|j r| rt d d } |dur|j durtd | |}|\}}| d ur|jstd|dd }d }|jr| s|durt|trt|tsd }t|t}n#t|tsd }t dt|}n|durttt}n|jsd}|dur|nd}| durtj||||jd} |durts||}tj|||jd}|jjr0|||| |dur+|jnd| }n|ddddddf}|j|jd}d|t|jj}|jru|duru|\}}}||f}|durotj||jd}||}nd}| ||jj!}| ||jj!}| rdnd}| rdnd}| r|jrdnd}d}d} |"|}!t#|j$D]\}"}#||"}$||"}%| r||!f}|j r|j r|%|#j&|!||||| |$|%d| | | }&n|#|!||||| |$|%|| | | d }&| d ur|&ddd|&dd}&|&dd\}!}'|&d}|jr|dur|&| rdnd} | r3||&df}|jr3||&df}q|'|!}!|"|!}!| rG||!f}| rL|'nd}(|rT|j}(|r[|(}(| slt)dd|!|(|||fDSt*|!|(|||dS) NZdecoder_zYou cannot specify both zinput_ids and zinputs_embeds at the same timer0zYou have to specify either z input_ids or inputs_embedszZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fztsz)Pop2PianoStack.forward..)last_hidden_statepast_key_valuesr6 attentionscross_attentions)+r>rroutput_hidden_statesuse_return_dictr\rsizerrorrerfrrNr r r Zfrom_legacy_cacheget_seq_lengthr%rrrr&_update_causal_maskrr1r4rr{Zinvert_attention_maskZ get_head_maskrrI enumeraterZ_gradient_checkpointing_funcr7rZto_legacy_cachetupler))r)rrrrr head_maskcross_attn_head_maskrrrrrrZerr_msg_prefixZ input_shaperrZreturn_legacy_cacheZreturn_self_attention_cachepast_key_values_lengthZmask_seq_lengthrZencoder_batch_sizeZencoder_sequence_length_Zencoder_hidden_shapeZencoder_extended_attention_maskZall_hidden_statesZall_attentionsZall_cross_attentionsrrr6rZ layer_modulerrZ layer_outputsZnext_decoder_cacheZ next_cacher-r-r.r7s2                 zPop2PianoStack.forwardFrr input_tensorrrrc Cs:|jjdkr|dur|dkr|SdS|jjdkr&t|tjr$t|}|S|dur.|nd}|dur7|jnd}|jjdkrO|sO|sOt j ||||j drOdS|j }|j d} |r^|} nt|tjri|j d n|| d} |j|| | |||j dd } |jjdkr|dur|jjd vr|st|j} t | | } | S) NZflash_attention_2rZflex_attentionrFZsdpa)rrZ is_trainingrr0)sequence_length target_lengthr4rr)cudaZxpuZnpu)r>Z_attn_implementationrrNr%rOrrZis_compileablerZ_ignore_causal_mask_sdparr4rZget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positionrtyperr{Z_unmask_unattended) r)rrrrrZpast_seen_tokensZusing_compilable_cacher4rrr min_dtyper-r-r.rsT           z"Pop2PianoStack._update_causal_maskrrr4rc KsD|dur|dkr|}|St|j}tj||f|||jd}|dkr+tj|dd}|tj||jd|ddk9}|ddddddf |ddd}|dur| }|j d} |ddddddd| f|ddddddf |j} | dk} |ddddddd| f | ||ddddddd| f<|S) aM Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. Nr)Z fill_valuer4rr)Zdiagonalrr0r)rrr%rr{rrZtriurZreshapeexpandrrr1Z masked_fill) rrrr4rrkwargsrrZ mask_lengthZ padding_maskr-r-r.rs,  $ 6  zDPop2PianoStack._prepare_4d_causal_attention_mask_with_cache_positionrM) NNNNNNNNNNNNN)F)r8r9r:r#rrr7rr%rOr rrrrr4rr;r-r-r+r.rsZ ]  Drcs(eZdZdZfddZddZZS)rz'Embedding Matrix for `composer` tokens.cs"ttj|j|jd|_dS)N)Znum_embeddingsZ embedding_dim)r"r#rrkcomposer_vocab_sizerCrrLr+r-r.r#s z&Pop2PianoConcatEmbeddingToMel.__init__cCs.||}||d}tj||gdd}|S)Nrrq)rrr%r)r)feature index_valueembedding_offsetZ index_shiftedZcomposer_embeddingrr-r-r.r7 sz%Pop2PianoConcatEmbeddingToMel.forward)r8r9r:__doc__r#r7r;r-r-r+r.rs rzA Pop2Piano Model with a `language modeling` head on top. )Z custom_introc*seZdZgdZdeffdd ZddZddZd d Zd d Z d dZ ddZ d2de j dededee j fddZe                  d3dee jdee j dee jdee jdee j dee j dee jdeeee jdeeee jd ee j dee j d!ee j d"ee jd#eed$eed%eed&eed'ee jd(eee j eff&d)d*Ze   + d4fd,d- Zd"e jfd.d/Zd0d1ZZS)5r)zencoder.embed_tokens.weightzdecoder.embed_tokens.weightzlm_head.weightr>cst|||_|j|_t|j|j|_t ||_ t |}d|_ d|_d|_t||j|_t |}d|_ d|_|j|_t||j|_tj|j|jdd|_|dS)NFTr@)r"r#r>rC model_dimrrkZ vocab_sizerrmel_conditionercopydeepcopyr\rZis_encoder_decoderrencoderZnum_decoder_layersrdecoderrBrr)r)r>Zencoder_configZdecoder_configr+r-r.r#s"     z*Pop2PianoForConditionalGeneration.__init__cCrrM)rrr-r-r.r6rz6Pop2PianoForConditionalGeneration.get_input_embeddingscCs"||_|j||j|dSrM)rrrrrr-r-r.r9s z6Pop2PianoForConditionalGeneration.set_input_embeddingscCrrMrrr-r-r.set_output_embeddings>rz7Pop2PianoForConditionalGeneration.set_output_embeddingscCrrMrrr-r-r.get_output_embeddingsArz7Pop2PianoForConditionalGeneration.get_output_embeddingscCrrM)rrr-r-r. get_encoderDrz-Pop2PianoForConditionalGeneration.get_encodercCrrM)rrr-r-r. get_decoderGrz-Pop2PianoForConditionalGeneration.get_decoderNinput_featurescomposergeneration_configrcCs|j}||vrtdt|d|||}tj||jd}||jd}t | }|j |||d}|durad||dddf <tj |dddfdd |gd d }||fS|dfS) a This method is used to concatenate mel conditioner tokens at the front of the input_features in order to control the type of MIDI token generated by the model. Args: input_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): input features extracted from the feature extractor. composer (`str`): composer token which determines the type of MIDI tokens to be generated. generation_config (`~generation.GenerationConfig`): The generation is used to get the composer-feature_token pair. attention_mask (``, *optional*): For batched generation `input_features` are padded to have the same shape across all examples. `attention_mask` helps to determine which areas were padded and which were not. - 1 for tokens that are **not padded**, - 0 for tokens that are **padded**. zPlease choose a composer from z. Composer received - rr)r r r Nrr0r)Zaxis)composer_to_feature_tokenkeysrrr%Ztensorrrepeatrr{rrrZ concatenater)r)rrrrrZcomposer_valuer r-r-r.get_mel_conditioner_outputsJs&  &z=Pop2PianoForConditionalGeneration.get_mel_conditioner_outputsrdecoder_input_idsdecoder_attention_maskrdecoder_head_maskrencoder_outputsrrdecoder_inputs_embedslabelsrrrrrreturncCs|dur|n|jj}|dur|n|jj}| dur | dur td| dur*| dur*| } |dur;|j||| ||||d}n$|r_t|ts_t|dt|dkrP|dndt|dkr[|dndd}|d}| durt|durt| durt|| }|j ||| | |||||||||d }|d}|jj r||j d }| |}d}| durt d d }||d |d | d }|s|f|dd|}|dur|f|S|St|||j|j|j|j|j|j|jd S)a` input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Pop2Piano is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining take a look a [Pop2Piano Training](./Pop2Piano#training). decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) Pop2Piano uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. input_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Does the same task as `inputs_embeds`. If `inputs_embeds` is not present but `input_features` is present then `input_features` will be considered as `inputs_embeds`. labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` NzSBoth `inputs_embeds` and `input_features` received! Please provide only one of them)rrrrrrrrrr/)rr6r) rrrrrrrrrrrrrrr)Z ignore_indexr0) lossZlogitsrZdecoder_hidden_statesZdecoder_attentionsrZencoder_last_hidden_staterZencoder_attentions)r>rrrrrNrrsrrrrrrrrrrr6rrr)r)rrr r!rr"rr#rrrr$r%rrrrrr6Zdecoder_outputsZsequence_outputZ lm_logitsr'Zloss_fctoutputr-r-r.r7{s|8    z)Pop2PianoForConditionalGeneration.forward composer1c s|dur|j}|jd i|t|dstdt|j|jjkr1td|jjdt|jd|j||||d\}}t j d d|||d|S) a Generates token ids for midi outputs. Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the model's default generation configuration. You can override any `generation_config` by passing the corresponding parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`. For an overview of generation strategies and code examples, check out the [following guide](./generation_strategies). Parameters: input_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): This is the featurized version of audio generated by `Pop2PianoFeatureExtractor`. attention_mask: For batched generation `input_features` are padded to have the same shape across all examples. `attention_mask` helps to determine which areas were padded and which were not. - 1 for tokens that are **not padded**, - 0 for tokens that are **padded**. composer (`str`, *optional*, defaults to `"composer1"`): This value is passed to `Pop2PianoConcatEmbeddingToMel` to generate different embeddings for each `"composer"`. Please make sure that the composet value is present in `composer_to_feature_token` in `generation_config`. For an example please see https://huggingface.co/sweetcocoa/pop2piano/blob/main/generation_config.json . generation_config (`~generation.GenerationConfig`, *optional*): The generation configuration to be used as base parametrization for the generation call. `**kwargs` passed to generate matching the attributes of `generation_config` will override them. If `generation_config` is not provided, the default will be used, which had the following loading priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s default values, whose documentation should be checked to parameterize generation. kwargs: Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. Return: [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True` or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`. Since Pop2Piano is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible [`~utils.ModelOutput`] types are: - [`~generation.GenerateEncoderDecoderOutput`], - [`~generation.GenerateBeamEncoderDecoderOutput`] Nrz`composer_to_feature_token` was not found! Please refer to https://huggingface.co/sweetcocoa/pop2piano/blob/main/generation_config.jsonand parse a dict like that.ztconfig.composer_vocab_size must be same as the number of keys in generation_config.composer_to_feature_token! Found z vs .)rrrr)Zinputsrrrr-) rrrrrsrr>r rr"generate)r)rrrrrr+r-r.r+s:6  z*Pop2PianoForConditionalGeneration.generatecCs ||SrM)r)r)r%r-r-r.%prepare_decoder_input_ids_from_labels]rzGPop2PianoForConditionalGeneration.prepare_decoder_input_ids_from_labelsc Cs|dur td|Sd}|D]M}d}|D]}||d||jf}q|dj|djkr@td|djd|djdt|t|krWtdt|dt|d||f}q|S) NzHYou might want to consider setting `use_cache=True` to speed up decodingr-rz%reordered_layer_past_states[0] shape z and layer_past_states[0] shape z mismatchedz&length of reordered_layer_past_states z! and length of layer_past_states )rewarningZ index_selectr1rrrrs)r)rZbeam_idxZreordered_decoder_pastZlayer_past_statesZreordered_layer_past_statesZlayer_past_stater-r-r._reorder_cache`s(  z0Pop2PianoForConditionalGeneration._reorder_cacherM)NNNNNNNNNNNNNNNNNN)Nr)N) r8r9r:Z_tied_weights_keysrr#rrrrrrr%Z FloatTensorstrrrrrZ LongTensorZ BoolTensorrOrrrrr7Zno_gradr+r,r.r;r-r-r+r.rs 1      Yr)Fr rr~typingrrrr%rZtorch.nnrZtransformers.generationrZ activationsr Z cache_utilsr r r Z generationr Zmodeling_attn_mask_utilsrZmodeling_outputsrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrrrZconfiguration_pop2pianorZ!torch.nn.attention.flex_attentionrZintegrations.flex_attentionrZ get_loggerr8reZ_load_pop2piano_layer_normZapex.normalizationrinfo ImportError Exceptionr-Moduler rr=rRrUrZrrrrrrr__all__r-r-r-r.sr                f%'fSu i