o iy@sdZddlZddlmZmZddlZddlmZddlmZddl m Z ddl m Z dd l mZmZmZdd lmZmZmZdd lmZdd lmZmZd dlmZmZmZmZmZd dl m!Z!d dl"m#Z#ddl$m%Z%erwddl&m'Z'm(Z(GdddeZ)eGdddeZ*GdddeZ+GdddeZ,eGddde*Z-eddGd d!d!e*e Z.Gd"d#d#eZ/Gd$d%d%e!Z0Gd&d'd'eZ1gd(Z2dS))zPyTorch PLBART model.N)OptionalUnion)nn)CrossEntropyLoss)Cache)GenerationMixin)AttentionMaskConverter_prepare_4d_attention_mask#_prepare_4d_attention_mask_for_sdpa)BaseModelOutputSeq2SeqLMOutputSeq2SeqModelOutput)PreTrainedModel)auto_docstringis_torch_flex_attn_available)BartClassificationHead BartDecoder BartEncoderBartForCausalLMBartScaledWordEmbedding)'BigBirdPegasusForSequenceClassification)shift_tokens_right) PLBartConfig) BlockMaskmake_flex_block_causal_maskc@ eZdZdS)PLBartScaledWordEmbeddingN__name__ __module__ __qualname__r$r$f/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/transformers/models/plbart/modular_plbart.pyr6rc @seZdZUeed<dZdZddgZdZdZ dZ de e j dfde j fd d Zdee e j d fd e j d e j defddZede j dedede jd e j def ddZde e j dfde e j dfde jde j fddZdS)PLBartPreTrainedModelconfigmodelTPLBartDecoderLayerPLBartEncoderLayerattention_maskN inputs_embedscCs|dur>|jjdkrd|vr|}|Sd}|S|jjdkr$t||j}|S|jjdkr8t|tjr6t|dd}|St||j}|S)Nflash_attention_2rsdpaflex_attentionF) is_causal r(_attn_implementationr dtype isinstancetorchTensorrr )selfr,r-r$r$r%_update_full_maskEs       z'PLBartPreTrainedModel._update_full_maskr input_tensorcache_positionpast_key_valuesc Csb|jjdkr*t|tjrt|}|S|dur(ttj|jd|jdf|jd}|S|jjdkr>|dur<|dk r<|SdS|durF| nd}|durO|j nd}|jjdkre|set j ||||jd redS|j}|jd}|rt|} nt|tjr|jd n||d} |j||| |||jdd } |jjdkr|dur|jjd vrt|j} t | | } | S) Nr0rr)sizedevicer.gFr/)r-past_key_values_length is_training)sequence_length target_lengthr4r; batch_size)cudaxpunpu)r(r3r5r6r7ronesshaper>anyget_seq_lengthis_compileabler _ignore_causal_mask_sdpatrainingr4get_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positiontypefinfomin_unmask_unattended) r8r,r:r;r<past_seen_tokensusing_compilable_cacher4rBrC causal_mask min_dtyper$r$r%_update_causal_mask\s`            z)PLBartPreTrainedModel._update_causal_maskrBrCr4rDc 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. N) fill_valuer4r>r)diagonalr>rAr)dimr6rRrSfullr>triuarangereshapeexpandclonerIto masked_fill) r,rBrCr4r;rDkwargsrWrX mask_length padding_maskr$r$r%rPs,  $ 6  zKPLBartPreTrainedModel._prepare_4d_causal_attention_mask_with_cache_positionencoder_hidden_statesencoder_attention_mask input_shapecCs|durM|durM|jjdkrd|vr|}|Sd}|S|jjdkr,t||j|dd}|S|jjdkrCt|tjrAt||ddd}|St||j|dd}|S) Nr.rr/rA)tgt_lenr0F) query_lengthr1r2)r8rjrkrlr-r$r$r%_update_cross_attn_masks2      z-PLBartPreTrainedModel._update_cross_attn_mask)r!r"r#r__annotations__base_model_prefixsupports_gradient_checkpointing_no_split_modules_supports_flash_attn_supports_sdpa_supports_flex_attnrr6r7r9rrrY staticmethodintr4rPSizeror$r$r$r%r':sZ    L 7  r'c@r) PLBartEncoderNr r$r$r$r%rzr&rzc@r) PLBartDecoderNr r$r$r$r%r{ r&r{c&s&eZdZddgZdeffdd ZddZdd Zd d Zd d Z e                d"de e j de e j de e j de e jde e jde e j de e jde ee jde ede e jde e jde ede ede ede ede e j deee jeff"d d!ZZS)# PLBartModelencoder.embed_tokens.weightdecoder.embed_tokens.weightr(cslt||j|j}}|jrt|jnd}t||j||d|_ t ||j |_ t ||j |_ |dS)Ng?) embed_scale)super__init__ pad_token_id vocab_sizescale_embeddingmathsqrtd_modelrsharedrzencoderr{decoder init_weights)r8r( padding_idxrr __class__r$r%rs  zPLBartModel.__init__cC|jSN)rr8r$r$r%get_input_embeddingsz PLBartModel.get_input_embeddingscCs||_|j|j_|j|j_dSr)rr embed_tokensr)r8valuer$r$r%set_input_embeddings s z PLBartModel.set_input_embeddingscCs4|jjr||jj|j||jj|jdSdSr)r(tie_word_embeddings_tie_or_clone_weightsrrrrrr$r$r% _tie_weights%szPLBartModel._tie_weightscCrr)rrr$r$r% get_encoder*rzPLBartModel.get_encoderN input_idsr,decoder_input_idsdecoder_attention_mask head_maskdecoder_head_maskcross_attn_head_maskencoder_outputsr<r-decoder_inputs_embeds use_cacheoutput_attentionsoutput_hidden_states return_dictr;returncCs4| 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`] or [`PLBartMultiTokenizer`] depending on the checkpoint. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) PLBart 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.* 50003 for *en_XX*, and 50001 for *java*. 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 (: obj:*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 (: obj:*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)rr,rr-rrrrrr)last_hidden_state hidden_states attentions) rr,rjrkrrr<r-rrrrr;)rr<decoder_hidden_statesdecoder_attentionscross_attentionsencoder_last_hidden_staterjencoder_attentions)r(rrruse_return_dictrrrr5r lenrrrr<rrr)r8rr,rrrrrrr<r-rrrrrr;decoder_outputsr$r$r%forward-sd1 zPLBartModel.forward)NNNNNNNNNNNNNNNN)r!r"r#_tied_weights_keysrrrrrrrrr6 LongTensorr7list FloatTensorrboolrtuplerr __classcell__r$r$rr%r| sv      r|zv The PLBART Model with a language modeling head. Can be used for code-to-text, text-to-code and code-to-code. ) custom_introc(sveZdZdZdgZgdZdeffdd ZddZd d Z d*d e de e de de jffdd Zd e dd fddZe d+de ejde ejde ejde ejde ejde ejde ejde eejde ede ejde ejd e ejd!e e d"e e d#e e d$e e d%e ejdeeejeff$d&d'Zd ejfd(d)ZZS),PLBartForConditionalGenerationr)final_logits_bias)r}r~zlm_head.weightr(csXt|t||_|dtd|jjjft j |j |jjjdd|_ | dS)NrrF)bias)rrr|r)register_bufferr6zerosrnum_embeddingsrLinearrlm_headr)r8r(rr$r%rs   z'PLBartForConditionalGeneration.__init__cC |jSr)r)rrr$r$r%r z*PLBartForConditionalGeneration.get_encodercCrr)r) get_decoderrr$r$r%rrz*PLBartForConditionalGeneration.get_decoderNTnew_num_tokenspad_to_multiple_of mean_resizingrcs&t|||}||jjd|S)Nr)rresize_token_embeddings_resize_final_logits_biasweightrI)r8rrrnew_embeddingsrr$r%rsz6PLBartForConditionalGeneration.resize_token_embeddingscCsj|jjd}||kr|jddd|f}ntjd||f|jjd}tj|j|gdd}|d|dS)NrArr])r^r)rrIr6rr>catr)r8rold_num_tokensnew_bias extra_biasr$r$r%rs z8PLBartForConditionalGeneration._resize_final_logits_biasrr,rrrrrrr<r-rlabelsrrrrr;cCs |dur|n|jj}| dur|dur| durt| |jj}|j|f|||||||| | | | ||||d}||d}||j|j}d}| dur[t }|| d|jj | d}|sq|f|dd}|duro|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`] or [`PLBartMultiTokenizer`] depending on the checkpoint. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) PLBart 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.* 50003 for *en_XX*, and 50001 for *java*. 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 (: obj:*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 (: obj:*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 Mask-filling: ```python >>> from transformers import AutoTokenizer, PLBartForConditionalGeneration >>> model = PLBartForConditionalGeneration.from_pretrained("uclanlp/plbart-base") >>> tokenizer = AutoTokenizer.from_pretrained("uclanlp/plbart-base") >>> # en_XX is the language symbol id for English >>> TXT = " Is 0 the Fibonacci number ? en_XX" >>> 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() ['first', 'same', 'highest', 'result', 'number'] ``` N)r,rrrrrrr<r-rrrrrr;rrAr) losslogitsr<rrrrrjr)r(rrrr)rrrer>rviewrr r<rrrrrjr)r8rr,rrrrrrr<r-rrrrrrr;outputs lm_logitsmasked_lm_lossloss_fctoutputr$r$r%rsVKz&PLBartForConditionalGeneration.forwardcCst||jjSr)rr(r)r8rr$r$r%%prepare_decoder_input_ids_from_labels?szDPLBartForConditionalGeneration.prepare_decoder_input_ids_from_labels)NT)NNNNNNNNNNNNNNNNN)r!r"r#rq_keys_to_ignore_on_load_missingrrrrrrxrrr Embeddingrrrr6rr7rrrrrr rrrr$r$rr%rs       zrc@r)PLBartClassificationHeadNr r$r$r$r%rCr&rcseZdZfddZZS)PLBartForSequenceClassificationc tjdi|dS)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`] or [`PLBartMultiTokenizer`] depending on the checkpoint. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) PLBart 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.* 50003 for *en_XX*, and 50001 for *java*. 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 (: obj:*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 (: obj:*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). Nr$rr super_kwargsrr$r%rHs!z'PLBartForSequenceClassification.forward)r!r"r#rrr$r$rr%rGsrcs eZdZefddZZS)PLBartForCausalLMc r)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, PLBartForCausalLM >>> tokenizer = AutoTokenizer.from_pretrained("uclanlp/plbart-base") >>> model = PLBartForCausalLM.from_pretrained("uclanlp/plbart-base", 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 ```Nr$rrrr$r%rmszPLBartForCausalLM.forward)r!r"r#rrrr$r$rr%rlsr)rrrr|r')3__doc__rtypingrrr6rtorch.nnr cache_utilsr generationrmodeling_attn_mask_utilsr r r modeling_outputsr r rmodeling_utilsrutilsrrbart.modeling_bartrrrrr(bigbird_pegasus.modeling_bigbird_pegasusrmbart.modeling_mbartrconfiguration_plbartrintegrations.flex_attentionrrrr'rzr{r|rrrr__all__r$r$r$r%sH        K#% !