o wi8@sXdZddlZddlZddlmZmZddlZddlmZddlm Z ddl m Z ddl m Z dd lmZmZmZdd lmZdd lmZdd lmZdd lmZmZmZddlmZddlmZm 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/6dYnwGdddej7Z8e0se2Z8Gdddej7Z9Gdddej7Z:Gd d!d!ej7Z;Gd"d#d#ej7ZGd(d)d)eZ?e"Gd*d+d+eZ@Gd,d-d-e@ZAGd.d/d/ej7ZBe"d0d1Gd2d3d3e@eZCd3d+gZDdS)4zPyTorch Pop2Piano model.N)OptionalUnion)nn)CrossEntropyLoss)GenerationConfig)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)AttentionMaskConverter)GradientCheckpointingLayer)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentionsSeq2SeqLMOutput)PreTrainedModel) 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)self hidden_sizeeps __class__m/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/transformers/models/pop2piano/modeling_pop2piano.pyr"As  zPop2PianoLayerNorm.__init__cCs\|tjdjddd}|t||j}|jjtj tj fvr)||jj}|j|S)NT)keepdim) tor$float32powmeanrsqrtr'r&dtypefloat16bfloat16)r( hidden_statesvariancer-r-r.forwardIs  zPop2PianoLayerNorm.forward)r )__name__ __module__ __qualname__r"r< __classcell__r-r-r+r.r@src*eZdZdeffdd ZddZZS)Pop2PianoDenseActDenseconfigcsTttj|j|jdd|_tj|j|jdd|_t|j |_ t |j |_ dSNFbias)r!r"rLineard_modeld_ffwiwoDropout dropout_ratedropoutr dense_act_fnactr(rCr+r-r.r"_s zPop2PianoDenseActDense.__init__cCsl||}||}||}t|jjtjr/|j|jjjkr/|jjjtj kr/| |jjj}||}|SN) rJrPrN isinstancerKr&r$Tensorr7int8r2)r(r:r-r-r.r<fs    zPop2PianoDenseActDense.forwardr=r>r?rr"r<r@r-r-r+r.rB^srBcrA)Pop2PianoDenseGatedActDenserCcsjttj|j|jdd|_tj|j|jdd|_tj|j|jdd|_t |j |_ t |j |_dSrD)r!r"rrGrHrIwi_0wi_1rKrLrMrNrrOrPrQr+r-r.r"vs z$Pop2PianoDenseGatedActDense.__init__cCsz|||}||}||}||}t|jjtjr6|j |jjj kr6|jjj tj kr6| |jjj }||}|SrR) rPrXrYrNrSrKr&r$rTr7rUr2)r(r: hidden_gelu hidden_linearr-r-r.r<~s   z#Pop2PianoDenseGatedActDense.forwardrVr-r-r+r.rWusrWcrA)Pop2PianoLayerFFrCcsJt|jrt||_nt||_t|j|jd|_ t |j |_ dS)Nr*)r!r" is_gated_actrWDenseReluDenserBrrHlayer_norm_epsilon layer_normrrLrMrNrQr+r-r.r"s   zPop2PianoLayerFF.__init__cCs&||}||}|||}|SrR)rar_rN)r(r:forwarded_statesr-r-r.r<s  zPop2PianoLayerFF.forwardrVr-r-r+r.r\s r\csleZdZ  ddedeeffdd ZddZedd d Z dddZ         dddZ Z S)Pop2PianoAttentionFNrC 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.FrE)r!r" is_decoderhas_relative_attention_biasrelative_attention_num_bucketsrelative_attention_max_distancerHd_kvkey_value_proj_dim num_headsn_headsrMrN inner_dimrdlogger warning_oncer,r=rrGqkvo Embeddingrelative_attention_biasset pruned_headsgradient_checkpointingr(rCrfrdr+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) lenrrlrjrwrrprqrrrsrmunion)r(headsindexr-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) r2r$longabsmin zeros_likelogfloatmath full_likewhere)relative_position bidirectional num_buckets max_distancerelative_buckets max_exactis_smallrelative_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)r7device)rrr)r/rrr) rur&rr$arangerr2rrergrhpermute unsqueeze) r( query_length key_lengthrcache_positioncontext_positionmemory_positionrrelative_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)rr7)rrrrz)ptraining)'shaperpviewrlrj transpose is_updatedgetrdcross_attention_cacheself_attention_cache key_cache value_cacherqrrupdater$matmulrfzerosrr7rxr requires_gradrrwr%listboolr functionalsoftmaxrtype_asrN contiguousrmrs)r(r:maskkey_value_states position_biaspast_key_valuelayer_head_maskr use_cacheoutput_attentionsr batch_size seq_lengthis_cross_attention query_statesrcurr_past_key_valuecurrent_states key_states value_statesscoresrreal_seq_length causal_maskposition_bias_masked attn_weights attn_outputoutputsr-r-r.r<sx           " &     zPop2PianoAttention.forwardFN)Trr)NN) NNNNNNFFN) r=r>r?rrintr"r staticmethodrrr<r@r-r-r+r.rcs,#  /rccs@eZdZddeeffdd Z       d ddZZS) Pop2PianoLayerSelfAttentionFNrdcs>tt|||d|_t|j|jd|_t |j |_ dS)Nrfrdr]) r!r"rc SelfAttentionrrHr`rarrLrMrNryr+r-r.r"s z$Pop2PianoLayerSelfAttention.__init__c CsL||} |j| |||||||d} ||| d}|f| dd} | S)N)rrrrrrrrr)rarrN) r(r:attention_maskrrrrrrnormed_hidden_statesattention_outputrr-r-r.r<s  z#Pop2PianoLayerSelfAttention.forwardr)NNNNFFNr=r>r?rrr"r<r@r-r-r+r.rs rcsBeZdZddeeffdd Z        d ddZZS) Pop2PianoLayerCrossAttentionNrdcs>tt|d|d|_t|j|jd|_t |j |_ dS)NFrr]) r!r"rcEncDecAttentionrrHr`rarrLrMrN)r(rCrdr+r-r.r"s z%Pop2PianoLayerCrossAttention.__init__Fc  CsP||} |j| |||||||| | d } ||| d} | f| dd}|S)N) rrrrrrrrrrr)rarrN)r(r:rrrrrrrrrrr layer_outputrr-r-r.r<s  z$Pop2PianoLayerCrossAttention.forwardrR)NNNNFNFNrr-r-r+r.rs rcsJeZdZd deeffdd Z            d ddZZS) Pop2PianoBlockFNrdcs`t|j|_t|_|jt|||d|jr&|jt||d|jt |dS)Nr)rd) r!r"rer ModuleListlayerappendrrr\ryr+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)rmaxrr0)rrrrrrrr) rr7r$r8risinfanyfinforclampre)r(r:rrencoder_hidden_statesencoder_attention_maskencoder_decoder_position_biasrcross_attn_layer_head_maskrrr return_dictrself_attention_outputsattention_outputs clamp_valuedo_cross_attentioncross_attention_outputsrr-r-r.r<sn           zPop2PianoBlock.forwardr) NNNNNNNNFFTNrr-r-r+r.rsrc@s@eZdZeZdZdZdZdZdZ dgZ dgZ ddZ dd Z d S) Pop2PianoPreTrainedModel transformerFTrrKcCs|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?)r5stdlm_headrFN)"rCinitializer_factorrSrr&datafill_Pop2PianoConcatEmbeddingToMel embeddingnormal_!Pop2PianoForConditionalGenerationsharedhasattrtie_word_embeddingsrrBrJrHrFzero_rKrIrWrXrYrcrirkrprqrrrsrfru)r(modulefactorrHrjrlr-r-r. _init_weightsHsR               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.rzr).rz1self.model.config.pad_token_id has to be defined.) rCdecoder_start_token_id pad_token_id ValueErrorrr$fullrcat new_zerosclone masked_fill_)r( input_idsrrshifted_input_idsr-r-r. _shift_rightvs    z%Pop2PianoPreTrainedModel._shift_rightN)r=r>r?r config_classbase_model_prefixis_parallelizablesupports_gradient_checkpointing_supports_cache_class_supports_static_cache_no_split_modules_keep_in_fp32_modulesrr r-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).0irCr-r. sz+Pop2PianoStack.__init__..r]F)r!r" embed_tokensrerrrange num_layersblockrrHr`final_layer_normrLrMrN post_initmodel_parallel device_maprx)r(rCrr+rr.r"s   zPop2PianoStack.__init__cC|jSrRrr(r-r-r.get_input_embeddingsz#Pop2PianoStack.get_input_embeddingscC ||_dSrRr$r(new_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]k\}"}#||"}$||"}%| r||!f}|#|!||||| |$|%|| | | d }&| d ur|&ddd|&dd}&|&dd\}!}'|&d}|jr|dur|&| rdnd} | r||&df}|jr||&df}q|%|!}!|"|!}!| r,||!f}| r1|'nd}(|r9|j}(|r@|&}(| sQt'dd|!|(|||fDSt(|!|(|||dS) Ndecoder_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`...Fzbsz)Pop2PianoStack.forward..)last_hidden_statepast_key_valuesr: attentionscross_attentions))rCrroutput_hidden_statesuse_return_dictrersizerrxrrnrorrSr r r from_legacy_cacheget_seq_lengthr$rrrr%_update_causal_maskrr2r7rrinvert_attention_mask get_head_maskrrN enumeraterrto_legacy_cachetupler))r(r rrrr/ head_maskcross_attn_head_maskr5rrr8rrerr_msg_prefix input_shaperrreturn_legacy_cachereturn_self_attention_cachepast_key_values_lengthmask_seq_lengthrencoder_batch_sizeencoder_sequence_length_encoder_hidden_shapeencoder_extended_attention_maskall_hidden_statesall_attentionsall_cross_attentionsrrr:r layer_modulerr layer_outputsnext_decoder_cache next_cacher-r-r.r<s                 zPop2PianoStack.forwardFrr input_tensorrr5rc 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) Nflash_attention_2rflex_attentionrFsdpa)r/rI is_trainingrr0)sequence_length target_lengthr7rr)cudaxpunpu)rC_attn_implementationrrSr$rTrr<is_compileabler _ignore_causal_mask_sdparr7rget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positionrtyperr_unmask_unattended) r(rrWrr5rpast_seen_tokensusing_compilable_cacher7r\r]r min_dtyper-r-r.r=vsT           z"Pop2PianoStack._update_causal_maskr\r]r7rc 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. Nr1) fill_valuer7rr)diagonalr0r0r)r{r$rrrrtriurreshapeexpandr rr2 masked_fill) rr\r]r7rrkwargsrrj mask_length padding_maskr-r-r.res,  $ 6  zDPop2PianoStack._prepare_4d_causal_attention_mask_with_cache_positionrR) NNNNNNNNNNNNN)F)r=r>r?r"r&r+r<rr$rTr rr=rrr7rer@r-r-r+r.rsZ L  Drcs(eZdZdZfddZddZZS)rz'Embedding Matrix for `composer` tokens.cs"ttj|j|jd|_dS)N)num_embeddings embedding_dim)r!r"rrtcomposer_vocab_sizerHrrQr+r-r.r"s z&Pop2PianoConcatEmbeddingToMel.__init__cCs.||}||d}tj||gdd}|S)Nrrz)rrr$r)r(feature index_valueembedding_offset index_shiftedcomposer_embeddingr/r-r-r.r<sz%Pop2PianoConcatEmbeddingToMel.forward)r=r>r?__doc__r"r<r@r-r-r+r.rs rzA Pop2Piano Model with a `language modeling` head on top. ) 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.weightrCcst|||_|j|_t|j|j|_t ||_ t |}d|_ d|_d|_t||j|_t |}d|_ d|_|j|_t||j|_tj|j|jdd|_|dS)NFTrE)r!r"rCrH model_dimrrt vocab_sizerrmel_conditionercopydeepcopyreris_encoder_decoderrencodernum_decoder_layersrdecoderrGrr )r(rCencoder_configdecoder_configr+r-r.r" s"     z*Pop2PianoForConditionalGeneration.__init__cCr#rR)rr%r-r-r.r&$r'z6Pop2PianoForConditionalGeneration.get_input_embeddingscCs"||_|j||j|dSrR)rrr+rr)r-r-r.r+'s z6Pop2PianoForConditionalGeneration.set_input_embeddingscCr(rRrr)r-r-r.set_output_embeddings,r,z7Pop2PianoForConditionalGeneration.set_output_embeddingscCr#rRrr%r-r-r.get_output_embeddings/r'z7Pop2PianoForConditionalGeneration.get_output_embeddingscCr#rR)rr%r-r-r. get_encoder2r'z-Pop2PianoForConditionalGeneration.get_encodercCr#rR)rr%r-r-r. get_decoder5r'z-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 - r0r)rwrxryNrr0r)axis)composer_to_feature_tokenkeysrrr$tensorrrepeatrrrrr concatenater)r(rrrrrcomposer_valueryr-r-r.get_mel_conditioner_outputs8s&  &z=Pop2PianoForConditionalGeneration.get_mel_conditioner_outputsr decoder_input_idsdecoder_attention_maskrCdecoder_head_maskrDencoder_outputsr5r/decoder_inputs_embedslabelsrrr8rrreturncCs|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)r rr/rCrr8rrrr/)r4r:r6) r rr/r5rrrCrDrrr8rrrr) ignore_indexr0) losslogitsr5decoder_hidden_statesdecoder_attentionsr7encoder_last_hidden_staterencoder_attentions)rCrr9rrrSrr|r rrr~rrrr:rr5r:r6r7r4)r(r rrrrCrrDrr5r/rrrrrr8rrr:decoder_outputssequence_output lm_logitsrloss_fctoutputr-r-r.r<is|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)inputsr/rrr-) rrrrr|rrCrvrr!generate)r(rrrrrqr+r-r.rs:6  z*Pop2PianoForConditionalGeneration.generatecCs ||SrR)r )r(rr-r-r.%prepare_decoder_input_ids_from_labelsKr,zGPop2PianoForConditionalGeneration.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 )rnwarning index_selectr2rrrr|)r(r5beam_idxreordered_decoder_pastlayer_past_statesreordered_layer_past_stateslayer_past_stater-r-r._reorder_cacheNs(  z0Pop2PianoForConditionalGeneration._reorder_cacherR)NNNNNNNNNNNNNNNNNN)NrN) r=r>r?_tied_weights_keysrr"r&r+rrrrr$ FloatTensorstrrrrr LongTensor BoolTensorrTrBrrrr<no_gradrrrr@r-r-r+r.rs 1      Yr)Er|rrtypingrrr$rtorch.nnrtransformers.generationr activationsr cache_utilsr r r generationr modeling_attn_mask_utilsr modeling_layersrmodeling_outputsrrrmodeling_utilsr pytorch_utilsrrutilsrrrrrconfiguration_pop2pianor!torch.nn.attention.flex_attentionrintegrations.flex_attentionr get_loggerr=rn_load_pop2piano_layer_normapex.normalizationrinfo ImportError ExceptionrModulerrBrWr\rcrrrrrrr__all__r-r-r-r.sr                f%'fSd i