o Zh@sdZddlZddlZddlmZmZmZmZddlZddl Zddlm Z ddl m Z m Z mZddlmZddlmZmZdd lmZdd lmZmZmZdd lmZmZdd lmZmZm Z m!Z!m"Z"m#Z#m$Z$dd l%m&Z&ddl'm(Z(m)Z)m*Z*m+Z+ddl,m-Z-e)rddl.m/Z/ddl0m1Z1erddlm2Z2e+3e4Z5dej6de7fddZ8Gddde j9Z:Gddde j9Z;Gddde j<Z=Gddde=Z>Gd d!d!e=Z?e=e?e>d"Z@Gd#d$d$e j<ZAGd%d&d&e j<ZBGd'd(d(e j<ZCe(Gd)d*d*e&ZDGd+d,d,eDZEGd-d.d.eDZFe(Gd/d0d0eDZGe(d1d2Gd3d4d4eDeZHe(d5d2Gd6d7d7eDZIe(Gd8d9d9eDZJGd:d;d;eDZKGdZMdS)?zPyTorch MBART model.N)ListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)CacheEncoderDecoderCache)GenerationMixin)AttentionMaskConverter_prepare_4d_attention_mask#_prepare_4d_attention_mask_for_sdpa)!flash_attn_supports_top_left_maskis_flash_attn_available)BaseModelOutput)BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentionsSeq2SeqLMOutputSeq2SeqModelOutput#Seq2SeqQuestionAnsweringModelOutputSeq2SeqSequenceClassifierOutput)PreTrainedModel)auto_docstringis_torch_flex_attn_availableis_torchdynamo_compilinglogging) MBartConfig) BlockMask)make_flex_block_causal_mask)_flash_attention_forward input_ids pad_token_idcCs|}|dur td||dk|||jdddd}|d|}|ddddf|ddddf<||dddf<|S)z Shift input ids one token to the right, and wrap the last non pad token (the token) Note that MBart does not have a single `decoder_start_token_id` in contrast to other Bart-like models. Nz1self.model.config.pad_token_id has to be defined.ir dimr)clone ValueErrorZ masked_fill_nesum unsqueezegathersqueeze)r%r&Zprev_output_tokensZ index_of_eosZdecoder_start_tokensr1W/var/www/auris/lib/python3.10/site-packages/transformers/models/mbart/modeling_mbart.pyshift_tokens_rightCs(r3csJeZdZdZdedeffdd Zd dejd ed ejffd d ZZ S)MBartLearnedPositionalEmbeddingzN This module learns positional embeddings up to a fixed maximum size. num_embeddings embedding_dimcsd|_t||j|dSN)offsetsuper__init__)selfr5r6 __class__r1r2r;]sz(MBartLearnedPositionalEmbedding.__init__rNr%past_key_values_length position_idscs\|dur |jdd\}}tj|||tj|jjd|d}n|d}t ||j S)z3`input_ids' shape is expected to be [bsz x seqlen].Nr8)dtypedevicer)r) shapetorcharangelongweightrBexpandr.r:forwardr9)r<r%r?r@bszZseq_lenr=r1r2rIcs z'MBartLearnedPositionalEmbedding.forwardrN) __name__ __module__ __qualname____doc__intr;rDTensorrI __classcell__r1r1r=r2r4Xs(r4c sLeZdZdZd dedededeeffdd Zd ej ffd d Z Z S) MBartScaledWordEmbeddingz\ This module overrides nn.Embeddings' forward by multiplying with embeddings scale. ?r5r6 padding_idx embed_scalecst|||||_dSN)r:r;rV)r<r5r6rUrVr=r1r2r;ws z!MBartScaledWordEmbedding.__init__r%cst||jSrW)r:rIrV)r<r%r=r1r2rI{sz MBartScaledWordEmbedding.forward)rT) rLrMrNrOrPrfloatr;rDrQrIrRr1r1r=r2rSrs$rScseZdZdZ      ddededed ed ed ed eed eeffdd Z      dde j dee j dee dee j dee j dedee j de e j ee j ee e j ffddZZS)MBartAttentionz=Multi-headed attention from 'Attention Is All You Need' paperFTN embed_dim num_headsdropout is_decoderbias is_causalconfig layer_idxc st||_||_||_|||_||_|j||jkr*td|jd|d|jd|_||_ ||_ ||_ |durK|j rKt d|jjdtj|||d|_tj|||d|_tj|||d|_tj|||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).gzInstantiating a decoder z without passing `layer_idx` is not recommended and will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.r_)r:r;r[r\r]head_dimrar+scalingr^r`rblogger warning_oncer>rLrLineark_projv_projq_projout_proj) r<r[r\r]r^r_r`rarbr=r1r2r;s0    zMBartAttention.__init__ hidden_stateskey_value_statespast_key_valueattention_masklayer_head_maskoutput_attentionscache_positionreturncCsv|du}|\} } } ||| d|j|jdd} | |j} |dur=t|tr;|j |j } |r7|j }n|j }n|}|rA|n|}|rX|durX| rX|j|j }|j|j }nE||}||}|| d|j|jdd}|| d|j|jdd}|dur|s|nd}||||j d|i\}}|rd|j |j <| |jd|jf}| j|} |j|}|j|}|d}t| |dd}|| |j| |fkrtd| |j| |fd||dur|ddddddd|jd f}|| |j| ||}|| |j| |}tjj|dd }|durN||jfkr3td |jfd||dddd|| |j| |}|| |j| |}|re|| |j| |}|| |j| |}nd}tjj||j|jd }t||}|| |j| |jfkrtd | |j| |jfd||| |j| |j}|dd}|| | |j}||}|||fS)#Input shape: Batch x Time x ChannelNr)r r8rsTz$Attention weights should be of size z , but is r'z/Head mask for a single layer should be of size ptrainingz `attn_output` should be of size )sizerkviewr\rd transposere isinstancer is_updatedgetrbcross_attention_cacheself_attention_cache key_cache value_cacherirjupdatereshaperDZbmmr+rCr functionalZsoftmaxr]ryr[rl)r<rmrnrorprqrrrsis_cross_attentionrJtgt_len_ query_statesr~curr_past_key_valuecurrent_states key_states value_statesZ proj_shapeZsrc_len attn_weightsZattn_weights_reshapedZ attn_probs attn_outputr1r1r2rIs"             &  "    zMBartAttention.forward)rZFTFNNNNNNFN)rLrMrNrOrPrXboolrr!r;rDrQr rrIrRr1r1r=r2rYs` * rYcseZdZdZfddZ      ddejdeejdeed eejd eejd e d eejd e ejeejee ejffddZ Z S)MBartFlashAttention2aF MBart flash attention module. This module inherits from `MBartAttention` as the weights of the module stays untouched. The only required change would be on the forward pass where it needs to correctly call the public API of flash attention and deal with padding tokens in case the input contains any of them. cstj|i|t|_dSrW)r:r;r_flash_attn_uses_top_left_maskr<argskwargsr=r1r2r;%s zMBartFlashAttention2.__init__NFrmrnrorprqrrrsrtc Cs|rtd|du}|\} } } ||| d|j|j} |dur:t|tr8|j |j } |r4|j }n|j }n|}|r>|n|}|rU|durU| rU|j |j }|j|j }nE||}||}|| d|j|jdd}|| d|j|jdd}|dur|s|nd}||||j d|i\}}|rd|j|j <|dd}|dd}| j}|tjkrtrt}nt|jdr|jj}n|jjj}td|d | |} ||}||}t| |||| |j r|j!nd |j"|j#d }|$| | d}|%|}|d|fS) NzMBartSdpaAttention2 attention does not support `output_attentions`. Use the argument `attn_implementation='eager'` when loading the model.r)r r8rsT_pre_quantization_dtypezThe input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in .rZ)r]r`Zuse_top_left_mask)&r+rzrkr{r\rdr}r r~rrbrrrrrirjr|rrArDZfloat32Zis_autocast_enabledZget_autocast_gpu_dtypehasattrrarrGrfrgtor$ryr]r`rrrl)r<rmrnrorprqrrrsrrJZq_lenrrr~rrrrZ input_dtypeZ target_dtyperr1r1r2rI-sv                    zMBartFlashAttention2.forwardr) rLrMrNrOr;rDrQrr rrrIrRr1r1r=r2rs4   rcseZdZ      d dejdeejdeedeejdeejded eejd eejeejeeejfffd d Z Z S)MBartSdpaAttentionNFrmrnrorprqrrrsrtcs,|rtdtj||||||dS|du}|\} } } ||| d|j|j dd} |durKt |t rI|j |j} |rE|j}n|j}n|}|rO|n|}|rf|durf| rf|j|j}|j|j}nE||}||}|| d|j|j dd}|| d|j|j dd}|dur|s|nd}||||jd|i\}}|rd|j |j<d}|dur|ddddddd|jd f}| jjd kr|dur| } |}|}|jr|dur| dkrdnd }tjjj| ||||jr|j nd |d }| dd}|| | |j!}|"|}|d|fS)ruaMBartModel is using MBartSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` . Falling back to the manual attention implementation, but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.)rnrorprrrsNr)r r8rsTrvcudaFrZ) attn_maskZ dropout_pr`)#rfrgr:rIrzrkr{r\rdr|r}r r~rrbrrrrrirjrrCrBtype contiguousr`rDrrZscaled_dot_product_attentionryr]r[rl)r<rmrnrorprqrrrsrrJrrrr~rrrr causal_maskr`rr=r1r2rIsr  "       &   zMBartSdpaAttention.forwardr) rLrMrNrDrQrr rrrIrRr1r1r=r2rs0 r)eagersdpaflash_attention_2c sLeZdZdeffdd Z d dejdejdejded ejf d d ZZ S) MBartEncoderLayerracst|j|_t|j|j|j|j|d|_t |j|_ |j |_ t |j|_|j|_t |j|j|_t |j|j|_t |j|_dS)N)r[r\r]ra)r:r;d_modelr[MBART_ATTENTION_CLASSES_attn_implementationZencoder_attention_headsattention_dropout self_attnr LayerNormself_attn_layer_normr]r activation_function activation_fnactivation_dropoutrhZencoder_ffn_dimfc1fc2final_layer_normr<rar=r1r2r;s  zMBartEncoderLayer.__init__Frmrprqrrrtc Cs|}||}|j||||d\}}}tjj||j|jd}||}|}||}|||}tjj||j |jd}| |}tjj||j|jd}||}|j t j krht |j jd}t j|| |d}|f} |rr| |f7} | S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. )rmrprqrrrwi)minmax)rrrrr]ryrrrrrrArDZfloat16finforclamp) r<rmrprqrrresidualrrZ clamp_valueoutputsr1r1r2rIs0      zMBartEncoderLayer.forwardF) rLrMrNr!r;rDrQrrIrRr1r1r=r2rsrcseZdZddedeeffdd Z         ddejd eejd eejd eejd eejd eejdee dee dee deejdejfddZ Z S)MBartDecoderLayerNrarbc st|j|_t|j|j|j|jdd||d|_|j |_ t |j |_ |j |_ t|j|_t|j|j|j|jd||d|_t|j|_t|j|j|_t|j|j|_t|j|_dS)NT)r[r\r]r^r`rarb)r]r^rarb)r:r;rr[rrZdecoder_attention_headsrrr]r rrrrrr encoder_attnencoder_attn_layer_normrhZdecoder_ffn_dimrrr)r<rarbr=r1r2r;Ks6  zMBartDecoderLayer.__init__FTrmrpencoder_hidden_statesencoder_attention_maskrqcross_attn_layer_head_maskrorr use_cachersrtc Cs |} ||}|j|||||| d\}} }tjj||j|jd}| |}d} |durN|} ||}|j||||||d\}} }tjj||j|jd}| |}|} ||}| | |}tjj||j |jd}| |}tjj||j|jd}| |}|f}|r|| | f7}| r||f7}|S)a8 Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(batch, seq_len, embed_dim)` encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of size `(decoder_attention_heads,)`. past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. It is used to update the cache in the correct position and to infer the complete sequence length. )rmrorprqrrrsrwN)rmrnrprqrorr) rrrrr]ryrrrrrrr)r<rmrprrrqrrorrrrsrZself_attn_weightsZcross_attn_weightsrr1r1r2rIjsN"        zMBartDecoderLayer.forwardrW) NNNNNNFTN) rLrMrNr!rrPr;rDrQr rrIrRr1r1r=r2rJsD"    rcsHeZdZdZdedededeffdd Zdejd ejfd d Z Z S) MBartClassificationHeadz-Head for sentence-level classification tasks. input_dim inner_dim num_classespooler_dropoutcs8tt|||_tj|d|_t|||_dS)N)rx)r:r;rrhdenseZDropoutr]rl)r<rrrrr=r1r2r;s z MBartClassificationHead.__init__rmrtcCs6||}||}t|}||}||}|SrW)r]rrDtanhrl)r<rmr1r1r2rIs     zMBartClassificationHead.forward) rLrMrNrOrPrXr;rDrQrIrRr1r1r=r2rs rc @seZdZeZdZdZgdZdZdZ dZ dZ ddZ e ddZ dd eejd fd ejd ejd edef ddZed ejdededejd ejdef ddZdS)MBartPreTrainedModelmodelT)rrrYcCs|jj}t|tjr"|jjjd|d|jdur |jj dSdSt|tj rC|jjjd|d|j durA|jj|j  dSdSt|tj rX|jj d|jj dSdS)NrZ)meanstdrT)raZinit_stdr}rrhrGdataZnormal_r_Zzero_ EmbeddingrUrZfill_)r<modulerr1r1r2 _init_weightss     z"MBartPreTrainedModel._init_weightscCs>|jj}tjgddddd|gg|jd}|||d}|S)N)r r8r r8rB)rpr%)rar&rDZtensorrBr,)r<Z pad_tokenr% dummy_inputsr1r1r2rs "z!MBartPreTrainedModel.dummy_inputsFrpr" input_tensorrspast_key_valuesrrc 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) NrrZZflex_attentionrFr) inputs_embedsr?Z is_trainingr r))sequence_length target_lengthrArs batch_size)rZxpuZnpu)raranyr}rDrQr#get_seq_lengthZis_compileablerZ_ignore_causal_mask_sdparyrArCZget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positionrBrrrZ_unmask_unattended) r<rprrsrrrZpast_seen_tokensZusing_compilable_cacherArrr min_dtyper1r1r2_update_causal_masksT           z(MBartPreTrainedModel._update_causal_maskrrrArc 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_valuerArBr )Zdiagonalrr)r)r(rDrrfullrBZtriurErrHr*rCrZ masked_fill) rprrrArsrrrrZ mask_lengthZ padding_maskr1r1r2rBs,  $ 6  zJMBartPreTrainedModel._prepare_4d_causal_attention_mask_with_cache_positionNr)rLrMrNr!Z config_classbase_model_prefixsupports_gradient_checkpointingZ_no_split_modulesZ_supports_flash_attn_2Z_supports_sdpaZ_supports_cache_classZ_supports_static_cacherpropertyrrrDrQr rr staticmethodrPrArr1r1r1r2rsL   DrcseZdZdZddedeejffdd ZddZ       dd ee j d ee j d ee j d ee j d eedeedeedeeeffddZZS) MBartEncoderz Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`MBartEncoderLayer`]. Args: config: MBartConfig embed_tokens (nn.Embedding): output embedding Nra embed_tokenscstj|_j|_j}j|_j|_ j r!t |nd}t j||j|d|_|dur7|j|j_tj||_tfddtjD|_|_t||_tj|_d|_|dS)NrTrVcsg|]}tqSr1)r).0rrar1r2 sz)MBartEncoder.__init__..F)r:r;r]Zencoder_layerdrop layerdroprr&rUmax_position_embeddingsZmax_source_positionsscale_embeddingmathsqrtrS vocab_sizerrGr4embed_positionsr ModuleListrangeZencoder_layerslayersrarlayernorm_embedding layer_normgradient_checkpointing post_init)r<rarr[rVr=rr2r;s,      zMBartEncoder.__init__cCs(|jrt|jddr|dSdSdS)NrF)rgetattrraZgradient_checkpointing_enabler<r1r1r2._backward_compatibility_gradient_checkpointings z;MBartEncoder._backward_compatibility_gradient_checkpointingr%rp head_maskrrroutput_hidden_states return_dictrtcCs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur*|dur*td|dur<|}|j} |d| d}n|durL|dddddf}ntd|durY||}||} || |j } | | } t j j| |j|jd} |dur|jjdkrd|vr|nd}n|jjdkr|dur|st||j}nt||j}|rd nd} |rd nd} |dur|dt|jkrtd t|jd |dd t|jD]\\}}|r| | f} d }|jrtg}||jkrd}|rd}n1|jr|jr||j| ||dur ||nd|}n|| ||dur||nd|d}|d} |r0| |df} q|| } |r>| | f} |sMtdd| | | fDSt | | | dS)a~ Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. NzDYou cannot specify both input_ids and inputs_embeds at the same timer)z5You have to specify either input_ids or inputs_embedsrwrrrr1z&The head_mask should be specified for  layers, but it is for rFT)NN)rqrrr cs|] }|dur|VqdSrWr1rvr1r1r2 1sz'MBartEncoder.forward..last_hidden_staterm attentions)!rarrruse_return_dictr+rCr{rrrrBrrrr]ryrrrArrzlenr enumeraterDrandrr_gradient_checkpointing_func__call__rtupler)r<r%rprrrrrrinputZ input_shapeZ embed_posrmZencoder_statesZall_attentionsidxZ encoder_layerZto_dropdropout_probability layer_outputsr1r1r2rIs.             zMBartEncoder.forwardrW)NNNNNNN)rLrMrNrOr!rrrr;rrD LongTensorrQ FloatTensorrrrrrIrRr1r1r=r2r{s8  rcseZdZdZddedeejffdd ZddZ d d Z             dd ee j d ee j d ee jdee j dee j dee j deeee jdee jdeedeedeedeedee j deeeffddZZS) MBartDecoderz Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`MBartDecoderLayer`] Args: config: MBartConfig embed_tokens (nn.Embedding): output embedding Nrarcstj|_j|_j|_j|_j rt j nd}t jj |j|d|_|dur6|j|j_tjj |_tfddtjD|_|_tj |_tj |_d|_|dS)NrTrcsg|]}t|dqS))rb)r)rirr1r2rSsz)MBartDecoder.__init__..F)r:r;r]Zdecoder_layerdroprr&rUrZmax_target_positionsrrrrrSrrrGr4rrrrZdecoder_layersrrarrrrr)r<rarrVr=rr2r;@s*    zMBartDecoder.__init__cC|jSrWrrr1r1r2get_input_embeddings]z!MBartDecoder.get_input_embeddingscC ||_dSrWrr<valuer1r1r2set_input_embeddings` z!MBartDecoder.set_input_embeddingsr%rprrrcross_attn_head_maskrrrrrrrrsrtc"CsN| dur| n|jj} | dur| n|jj} | dur| n|jj} | dur$| n|jj} |du|duAr4td|durA|d|jd}|durJ||}|j rY|j rY| rYt dd} d}| rnt |tsnd}t dt|}|dd\}}|dur|nd}| durtj||||jd } |durts||}tj|||jd }t |tr|jn|}|||| || }|dur|dur|jjd krd|vr|nd}n|jjd kr|dur| st||j|d }nt||j|d }|jt|| d }|| |j}|!|}t"j#j$||j$|j d}| rdnd}| rdnd}| r$|dur$dnd}d}t%||gddgD]+\}}|dur[|dt&|j'kr[td|dt&|j'd|ddq1t(|j'D]\}}| rn||f7}|j rt)g}||j*krqb|j r|j r|+|j,|||||dur||nd|dur||ndd| | | } n ||||||dur||nd|dur||nd|| | | d } | d}| r| | rdnd}| r|| df7}|dur|| df7}qb|-|}| r||f7}| r|nd}!|r |.}!| st/dd||!|||fDSt0||!|||dS)a4 Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*): Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. 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 `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing cross-attention on hidden heads. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. It is used to update the cache in the correct position and to infer the complete sequence length. NzTYou cannot specify both decoder_input_ids and decoder_inputs_embeds at the same timer)z[`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`...FTzPassing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. You should pass an instance of `EncoderDecoderCache` instead, e.g. `past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`.rrrr)r)r@rwr1rr"zThe `z` should be specified for rr) rprrrqrrorrrrsr r r8csrrWr1rr1r1r2rXsz'MBartDecoder.forward..)rrrmr cross_attentions)1rarrrrr r+r{rCrrryrfrgr}r r Zfrom_legacy_cacherzrrDrErBrZonesrrrrrArrrrrrrr]zipr rr r rrrrZto_legacy_cacherr)"r<r%rprrrr"rrrrrrrrsZreturn_legacy_cacherZ seq_lengthr?Zmask_seq_lengthZself_attn_cacherr@rmZall_hidden_statesZall_self_attnsZall_cross_attentionsZnext_decoder_cacherZ mask_namerZ decoder_layerrrZ next_cacher1r1r2rIcsT               zMBartDecoder.forwardrW) NNNNNNNNNNNNN)rLrMrNrOr!rrrr;rr rDrrQrrrrrrIrRr1r1r=r2r7s^      rc's<eZdZddgZdeffdd ZddZdd Zd d Zd d Z ddZ e                d$de e jde e jde e jde e jde e jde e jde e jde eee jde eee jde e jde e jde ede ede ede ed e e jd!eeee jff"d"d#ZZS)% MBartModelzencoder.embed_tokens.weightzdecoder.embed_tokens.weightracslt||j|j}}|jrt|jnd}t||j||d|_ t ||j |_ t ||j |_ |dS)NrTr)r:r;r&rrrrrrSsharedrencoderrdecoderr)r<rarUrrVr=r1r2r;js  zMBartModel.__init__cCrrW)r&rr1r1r2rwrzMBartModel.get_input_embeddingscCs||_|j|j_|j|j_dSrW)r&r'rr(rr1r1r2r zs zMBartModel.set_input_embeddingscCrrW)r'rr1r1r2 get_encoderrzMBartModel.get_encodercCrrWr(rr1r1r2 get_decoderrzMBartModel.get_decodercCs8|jjr||jj|||jj|dSdSrW)raZtie_word_embeddingsZ_tie_or_clone_weightsr'rrr(rr1r1r2 _tie_weightsszMBartModel._tie_weightsNr%rpdecoder_input_idsdecoder_attention_maskrdecoder_head_maskr"encoder_outputsrrdecoder_inputs_embedsrrrrrrsrtcCs4| dur| n|jj} |dur|n|jj}| dur| n|jj} |dur$|n|jj}|dur7| dur7t||jj}|durH|j|||| | ||d}n$|rlt|t slt |dt |dkr]|dndt |dkrh|dndd}|j |||d|||| | | | |||d }|s||St |j |j|j|j|j|j |j|jdS) a 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) MBart uses a specific language id token as the starting token for `decoder_input_ids` generation that varies according to source and target language, *e.g.* 25004 for *en_XX*, and 25003 for *de_DE*. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). For translation and summarization training, `decoder_input_ids` should be provided. If no `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right for denoising pre-training following the paper. decoder_attention_mask (`torch.LongTensor` 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. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_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**. N)r%rprrrrrrrr r8r r%rprrrr"rrrrrrrrs)rrdecoder_hidden_statesdecoder_attentionsr#encoder_last_hidden_staterencoder_attentions)rarrrrr r3r&r'r}rr r(rrrrmr r#)r<r%rpr-r.rr/r"r0rrr1rrrrrrsZdecoder_outputsr1r1r2rIsd/ zMBartModel.forwardNNNNNNNNNNNNNNNN)rLrMrN_tied_weights_keysr!r;rr r)r+r,rrrDrrQrrrrrrIrRr1r1r=r2r%fsx      r%z The MBART Model with a language modeling head. Can be used for summarization, after fine-tuning the pretrained models. )Z custom_introc)seZdZdZdgZgdZdeffdd ZddZd d Z d0d e de e de de jffdd Zd e dd fddZddZddZe d1de ejde ejde ejde ejde ejde ejde ejd e eeejd!e eeejd"e ejd#e ejd$e ejd%e e d&e e d'e e d(e e d)e ejdeeeejff$d*d+Zd$ejfd,d-Zed.d/ZZS)2MBartForConditionalGenerationrfinal_logits_bias)!model.encoder.embed_tokens.weight!model.decoder.embed_tokens.weightlm_head.weightracsXt|t||_|dtd|jjjft j |j |jjjdd|_ | dS)Nr:r Frc)r:r;r%rregister_bufferrDzerosr&r5rrhrlm_headrrr=r1r2r;s   z&MBartForConditionalGeneration.__init__cC |jSrW)rr)rr1r1r2r) r!z)MBartForConditionalGeneration.get_encodercCrArW)rr+rr1r1r2r+ r!z)MBartForConditionalGeneration.get_decoderNTnew_num_tokenspad_to_multiple_of mean_resizingrtcs&t|||}||jjd|S)Nr)r:resize_token_embeddings_resize_final_logits_biasrGrC)r<rBrCrDnew_embeddingsr=r1r2rEsz5MBartForConditionalGeneration.resize_token_embeddingscCsj|jjd}||kr|jddd|f}ntjd||f|jjd}tj|j|gdd}|d|dS)Nr)r rr'r:)r:rCrDr?rBcatr>)r<rBZold_num_tokensZnew_biasZ extra_biasr1r1r2rFs z7MBartForConditionalGeneration._resize_final_logits_biascCrrWr@rr1r1r2get_output_embeddings rz3MBartForConditionalGeneration.get_output_embeddingscCrrWrIr<rGr1r1r2set_output_embeddings#r!z3MBartForConditionalGeneration.set_output_embeddingsr%rpr-r.rr/r"r0rrr1labelsrrrrrrscCs|dur|n|jj}| dur&| rtdd} |dur&| dur&t| |jj}|j|f|||||||| | | | ||||d}||d|j}d}| dur^t }|| d|jj | d}|st|f|dd}|durr|f|S|St |||j |j|j|j|j|j|jd S) u 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) MBart uses a specific language id token as the starting token for `decoder_input_ids` generation that varies according to source and target language, *e.g.* 25004 for *en_XX*, and 25003 for *de_DE*. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). For translation and summarization training, `decoder_input_ids` should be provided. If no `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right for denoising pre-training following the paper. decoder_attention_mask (`torch.LongTensor` 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. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_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**. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example Translation: ```python >>> from transformers import AutoTokenizer, MBartForConditionalGeneration >>> model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-en-ro") >>> tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-en-ro") >>> example_english_phrase = "42 is the answer" >>> inputs = tokenizer(example_english_phrase, return_tensors="pt") >>> # Translate >>> generated_ids = model.generate(**inputs, num_beams=4, max_length=5) >>> tokenizer.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] '42 este răspuns' ``` Mask filling example: ```python >>> from transformers import AutoTokenizer, MBartForConditionalGeneration >>> model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-cc25") >>> tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-cc25") >>> # de_DE is the language symbol id for German >>> TXT = " Meine Freunde sind nett aber sie essen zu viel Kuchen. de_DE" >>> input_ids = tokenizer([TXT], add_special_tokens=False, return_tensors="pt")["input_ids"] >>> logits = model(input_ids).logits >>> masked_index = (input_ids[0] == tokenizer.mask_token_id).nonzero().item() >>> probs = logits[0, masked_index].softmax(dim=0) >>> values, predictions = probs.topk(5) >>> tokenizer.decode(predictions).split() ['nett', 'sehr', 'ganz', 'nicht', 'so'] ``` NzJThe `use_cache` argument is changed to `False` since `labels` is provided.F)rpr-r0r.rr/r"rrr1rrrrrrsrr)r losslogitsrr3r4r#r5rr6)rar rfwarningr3r&rr@r:rr{rrrr3r4r#r5rr6)r<r%rpr-r.rr/r"r0rrr1rMrrrrrrsrZ lm_logitsZmasked_lm_lossloss_fctoutputr1r1r2rI&sZ[ z%MBartForConditionalGeneration.forwardcCst||jjSrW)r3rar&)r<rMr1r1r2%prepare_decoder_input_ids_from_labelszCMBartForConditionalGeneration.prepare_decoder_input_ids_from_labelscsBd}|D]}|tfdd|ddD|ddf7}q|S)Nr1c3$|] }|d|jVqdSrKZ index_selectrrBrZ past_statebeam_idxr1r2r"z?MBartForConditionalGeneration._reorder_cache..r8rrrZZreordered_pastZ layer_pastr1rYr2_reorder_caches z,MBartForConditionalGeneration._reorder_cache)NTNNNNNNNNNNNNNNNNN) rLrMrNrZ_keys_to_ignore_on_load_missingr8r!r;r)r+rPrrrrrErFrJrLrrDrrQrrrrrIrTrr^rRr1r1r=r2r9s         r9z MBart model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. c&seZdZddgZdeffdd Ze                ddeej deej d eej d eej d eej d eej d eej dee ej deej deej deej dee dee dee dee deej deeeff"ddZZS)MBartForSequenceClassificationr;r<rac sBtj|fi|t||_t|j|j|j|j|_| dSrW) r:r;r%rrr num_labelsZclassifier_dropoutclassification_headr)r<rarr=r1r2r;s  z'MBartForSequenceClassification.__init__Nr%rpr-r.rr/r"r0rr1rMrrrrrrsrtcCs<|dur|n|jj}| durd} |dur!| dur!td|jj|j||||||||| | | | |||d}|d}||jj|j }t t | ddkrTtd||ddf|dd|ddddddf}||}d}| dur| |j } |jjdur|jjdkrd |j_n|jjdkr| jt jks| jt jkrd |j_nd |j_|jjd krt}|jjdkr||| }n,||| }n&|jjd krt}||d|jj| d}n|jjd krt}||| }|s |f|dd}|dur|f|S|St|||j|j|j|j|j |j!|j"d S) aN 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) Bart uses the `eos_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`). For translation and summarization training, `decoder_input_ids` should be provided. If no `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right for denoising pre-training following the paper. decoder_attention_mask (`torch.LongTensor` 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. If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_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**. labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). NFz8Passing input embeddings is currently not supported for rpr-r.rr/r"r0rr1rrrrrrsrr z7All examples must have the same number of tokens.r)Z regressionZsingle_label_classificationZmulti_label_classificationrN)#rar NotImplementedErrorr>rLreqZ eos_token_idrrBr rDZunique_consecutiver-r+r{rzrbZ problem_typerarArFrPr r0rrrrr3r4r#r5rr6)r<r%rpr-r.rr/r"r0rr1rMrrrrrrsrrmZeos_maskZsentence_representationrPrOrRrSr1r1r2rIs5 $     $       z&MBartForSequenceClassification.forwardr7)rLrMrNr8r!r;rrrDrrQrrrrrrrIrRr1r1r=r2r`sn       r`c(seZdZddgZfddZe                 ddeejdeejdeej d eej d eejd eejd eejd ee ej deej deej deej deej dee dee dee dee deej de eeff$ddZZS)MBartForQuestionAnsweringr;r<csBt|d|_|j|_t||_t|j|j|_| dSr7) r:r;rar%rrrh hidden_size qa_outputsrrr=r1r2r;^s   z"MBartForQuestionAnswering.__init__Nr%rpr-r.rr/r"r0start_positions end_positionsrr1rrrrrrsrtcCs||dur|n|jj}| dur| durd} |j||||||||| | | ||||d}|d}||}|jddd\}}|d}|d}d}| dur| durt| dkr_| d} t| dkrl| d} |d}| d|} | d|} t |d}||| }||| }||d }|s||f|dd}|dur|f|S|St ||||j |j |j|j|j|j|jd S) a$ 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) Bart uses the `eos_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`). For translation and summarization training, `decoder_input_ids` should be provided. If no `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right for denoising pre-training following the paper. decoder_attention_mask (`torch.LongTensor` 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. If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_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**. NFrcrr r)r')Z ignore_indexr8) rO start_logits end_logitsrr3r4r#r5rr6)rar rrhsplitr0rr rzrrrrr3r4r#r5rr6)r<r%rpr-r.rr/r"r0rirjrr1rrrrrrsrZsequence_outputrPrkrlZ total_lossZ ignored_indexrRZ start_lossZend_lossrSr1r1r2rIjsr3           z!MBartForQuestionAnswering.forwardr_)rLrMrNr8r;rrrDrQrrrrrrrrIrRr1r1r=r2rfZst        rfcs(eZdZdZfddZddZZS)MBartDecoderWrapperz This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is used in combination with the [`EncoderDecoderModel`] framework. cst|t||_dSrW)r:r;rr(rr=r1r2r;s zMBartDecoderWrapper.__init__cOs|j|i|SrWr*rr1r1r2rIszMBartDecoderWrapper.forward)rLrMrNrOr;rIrRr1r1r=r2rns rnc"seZdZdgZfddZddZddZdd Zd d Zd d Z ddZ e              d$de e jde e jde e jde e jde e jde e jde ee jde e jde e jde ede ede ede ede e jdeeeffd d!Zed"d#ZZS)%MBartForCausalLMr=csNt|}d|_d|_t|t||_tj |j |j dd|_ | dS)NTFrc)copydeepcopyr^Zis_encoder_decoderr:r;rnrrrhrgrr@rrr=r1r2r;s    zMBartForCausalLM.__init__cCs |jjjSrWrr(rrr1r1r2rr!z%MBartForCausalLM.get_input_embeddingscCs||jj_dSrWrrrr1r1r2r rUz%MBartForCausalLM.set_input_embeddingscCrrWrIrr1r1r2rJrz&MBartForCausalLM.get_output_embeddingscCrrWrIrKr1r1r2rLr!z&MBartForCausalLM.set_output_embeddingscCs ||j_dSrWrr()r<r(r1r1r2 set_decoder s zMBartForCausalLM.set_decodercCs|jjSrWrsrr1r1r2r+szMBartForCausalLM.get_decoderNr%rprrrr"rrrMrrrrrrsrtcCs| dur| n|jj} | dur| n|jj} | dur| n|jj} |jj||||||||| | | | |d }||d}d}| durU| |j} t }|| d|jj | d}| sk|f|dd}|duri|f|S|St |||j |j|j|jdS)a cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example: ```python >>> from transformers import AutoTokenizer, MBartForCausalLM >>> tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-cc25") >>> model = MBartForCausalLM.from_pretrained("facebook/mbart-large-cc25", add_cross_attention=False) >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder." >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits >>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size] >>> list(logits.shape) == expected_shape True ```Nr2rr)r )rOrPrrmr r#)rarrrr rr(r@rrBrr{rrrrmr r#)r<r%rprrrr"rrrMrrrrrrsrrPrOrRrSr1r1r2rIsH. zMBartForCausalLM.forwardcs.d}|D]}|tfdd|Df7}q|S)Nr1c3rVrKrWrXrYr1r2rpr[z2MBartForCausalLM._reorder_cache..r\r]r1rYr2r^ks zMBartForCausalLM._reorder_cache)NNNNNNNNNNNNNN)rLrMrNr8r;rr rJrLrtr+rrrDrrQrrrrrrrIrr^rRr1r1r=r2rosr          Yro)ror9rfr`r%r)NrOrprtypingrrrrrDZtorch.utils.checkpointrZtorch.nnrrr Z activationsr Z cache_utilsr r Z generationrZmodeling_attn_mask_utilsrrrZmodeling_flash_attention_utilsrrZmodeling_outputsrrrrrrrZmodeling_utilsrutilsrrrrZconfiguration_mbartr!Z!torch.nn.attention.flex_attentionr"Zintegrations.flex_attentionr#r$Z get_loggerrLrfrQrPr3rr4rSModulerYrrrrrrrrrr%r9r`rfrnro__all__r1r1r1r2s   $      xjDx!=1G