o i@sddlZddlmZmZddlZddlmZmZmZddlm Z ddl m Z ddl m Z mZmZddlmZdd lmZdd lmZmZmZdd lmZdd lmZmZmZdd lmZddl m!Z!ddl"m#Z#e$e%Z&Gdddej'Z(Gdddej'Z)Gdddej'Z*Gdddej'Z+Gdddej'Z,Gdddej'Z-GdddeZ.Gdd d ej'Z/Gd!d"d"ej'Z0Gd#d$d$ej'Z1Gd%d&d&ej'Z2Gd'd(d(ej'Z3Gd)d*d*eZ4Gd+d,d,e4Z5Gd-d.d.e4eZ6gd/Z7dS)0N)OptionalUnion)Tensordevicenn)CrossEntropyLoss)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)GradientCheckpointingLayer))BaseModelOutputWithPastAndCrossAttentions,BaseModelOutputWithPoolingAndCrossAttentions!CausalLMOutputWithCrossAttentions)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)logging)deprecate_kwarg)BlipTextConfigc s\eZdZdZfddZ    d deejdeejdeejd e d ej f d d Z Z S)BlipTextEmbeddingsz;Construct the embeddings from word and position embeddings.csttj|j|j|jd|_t|j|j|_ tj |j|j d|_ t |j |_|jdt|jdddt|dd|_||_dS) N) padding_idxeps position_ids)rF) persistentposition_embedding_typeabsolute)super__init__r Embedding vocab_size hidden_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutregister_buffertorcharangeexpandgetattrr!configselfr6 __class___/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/blip/modeling_blip_text.pyr$/s  zBlipTextEmbeddings.__init__Nr input_idsr inputs_embedspast_key_values_lengthreturnc Cs|dur |}n|dd}|d}|dur&|jdd|||f}|dur/||}|}|jdkr?||}||7}||}||}|S)Nrrr")sizerr)r!r+r,r0) r8r=rr>r? input_shape seq_length embeddingsr+r;r;r<forwardAs      zBlipTextEmbeddings.forward)NNNr) __name__ __module__ __qualname____doc__r$rr2 LongTensor FloatTensorintrrE __classcell__r;r;r9r<r,s$ rcseZdZdfdd ZddZddZdd Zd d Zed d dd       dde j de e j de e j de e j de e j d e e de ede e j dee j fddZZS)BlipTextSelfAttentionNcs&t||_|j|jdkrt|dstd|j|jf|j|_t|j|j|_|j|j|_ ||_ t |j|j |_ |rTt |j|j |_t |j|j |_nt |j|j |_t |j|j |_t |j|_t|dd|_|jdks~|jdkr|j|_t d|jd |j|_dSdS) Nrembedding_sizezLThe hidden size (%d) is not a multiple of the number of attention heads (%d)r!r" relative_keyrelative_key_queryr)r#r$r6r'num_attention_headshasattr ValueErrorrLattention_head_size all_head_size layer_idxrLinearqueryencoder_hidden_sizekeyvaluer.attention_probs_dropout_probr0r5r!r*r%distance_embeddingr8r6is_cross_attentionrXr9r;r<r$as0  zBlipTextSelfAttention.__init__cC ||_dSNattn_gradients)r8rer;r;r<save_attn_gradients} z)BlipTextSelfAttention.save_attn_gradientscC|jSrcrdr8r;r;r<get_attn_gradientsz(BlipTextSelfAttention.get_attn_gradientscCrbrc attention_map)r8rmr;r;r<save_attention_maprgz(BlipTextSelfAttention.save_attention_mapcCrhrcrlrir;r;r<get_attention_maprkz'BlipTextSelfAttention.get_attention_mappast_key_valuepast_key_values4.58new_nameversionF hidden_statesattention_mask head_maskencoder_hidden_statesencoder_attention_maskoutput_attentionscache_positionr@c Cs|j\} } } ||| d|j|jdd} |du} | r|n|}d}|dur?t|tr=|j |j }| r9|j }n|j }n|}| rC|n|}| r\|dur\|r\|j |j j}|j |j j}nF||| d|j|jdd}||| d|j|jdd}|dur| s|nd}||||j d|i\}}| rt|trd|j|j <t| |dd}|jdks|jd kr|d} tj| tj|jd dd}tj| tj|jd dd}||}|||jd}|j| jd }|jdkrtd | |}||}n|jd krtd | |}td ||}|||}|t |j}|dur.|||j}t!j"dd|}|#|}|durD||}t||}|$dddd%}|dd|j&f}|j|}||fS)NrrrRFr|TrPrQ)dtyperr~zbhld,lrd->bhlrzbhrd,lrd->bhlrdimrr)'shaperZviewrSrV transpose isinstancer is_updatedgetrXcross_attention_cacheself_attention_cachelayerskeysvaluesr\r]updater2matmulr!rAr3longrr_r*tor~einsummathsqrtrSoftmaxr0permute contiguousrW)r8rvrwrxryrzrqr{r| batch_sizerC_ query_layerrarcurr_past_key_valuecurrent_states key_layer value_layerattention_scoresposition_ids_lposition_ids_rdistancepositional_embeddingrelative_position_scoresrelative_position_scores_queryrelative_position_scores_keyattention_probsattention_probs_dropped context_layernew_context_layer_shaper;r;r<rEsz                zBlipTextSelfAttention.forwardrcNNNNNFN)rFrGrHr$rfrjrnrorr2rrrKr booltuplerErMr;r;r9r<rN`sB   rNc8eZdZfddZdejdejdejfddZZS)BlipTextSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNr) r#r$rrYr'denser,r-r.r/r0r7r9r;r<r$ zBlipTextSelfOutput.__init__rv input_tensorr@cC&||}||}|||}|Srcrr0r,r8rvrr;r;r<rE  zBlipTextSelfOutput.forwardrFrGrHr$r2rrErMr;r;r9r<r $rcseZdZdfdd ZddZeddd d       dd ejd eej d eej deej dee dee deejde ejfddZ ZS)BlipTextAttentionFNcs0tt|||d|_t||_t|_dS)NrX)r#r$rNr8routputset pruned_headsr`r9r;r<r$s   zBlipTextAttention.__init__cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrr)lenrr8rSrVrrrZr\r]rrrWunion)r8headsindexr;r;r< prune_heads s zBlipTextAttention.prune_headsrprqrrrsrvrwrxryr{r|r@c Cs>|j|||||||d}||d|} | f|dd} | S)Nrwrxryrqr{r|rr)r8r) r8rvrwrxryrqr{r| self_outputsattention_outputoutputsr;r;r<rEs  zBlipTextAttention.forward)FN)NNNNFN)rFrGrHr$rrr2rrrKr rrrErMr;r;r9r<rs6  rc2eZdZfddZdejdejfddZZS)BlipTextIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSrc) r#r$rrYr'intermediate_sizerr hidden_actstrr intermediate_act_fnr7r9r;r<r$9s   zBlipTextIntermediate.__init__rvr@cC||}||}|Src)rrr8rvr;r;r<rEA  zBlipTextIntermediate.forwardrr;r;r9r<r8s rcr)BlipTextOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) r#r$rrYrr'rr,r-r.r/r0r7r9r;r<r$IrzBlipTextOutput.__init__rvrr@cCrrcrrr;r;r<rEOrzBlipTextOutput.forwardrr;r;r9r<rHrrcseZdZfddZedddd       dd ejd eejd eejd eejd eejdee dee deejde ejfddZ ddZ ZS) BlipTextLayercsft||_|j|_d|_t||d|_||_|jjr't||jj|d|_ t ||_ t ||_ dS)Nrr)rarX)r#r$r6chunk_size_feed_forward seq_len_dimr attention layer_num is_decodercrossattentionr intermediaterr)r8r6rr9r;r<r$Ws   zBlipTextLayer.__init__rprqrrrsNFrvrwrxryrzr{r|r@c  Cs|j||||||d} | d} | dd} |dur1|j| ||||||d} | d} | | dd} t|j|j|j| } | f| S)N)rwrxr{rqr|rrr)rrrfeed_forward_chunkrr)r8rvrwrxryrzrqr{r|self_attention_outputsrrcross_attention_outputs layer_outputr;r;r<rEes4    zBlipTextLayer.forwardcCs||}|||}|Src)rr)r8rintermediate_outputrr;r;r<rs  z BlipTextLayer.feed_forward_chunkr)rFrGrHr$rr2rrrKr rrrErrMr;r;r9r<rVs<    'rcseZdZfddZ          ddejdeejdeejd eejd eejd eed ee d ee dee dee deejde e eje ffddZ ZS)BlipTextEncodercs:t|_tfddtjD|_d|_dS)Ncsg|]}t|qSr;)r).0ir6r;r< z,BlipTextEncoder.__init__..F) r#r$r6r ModuleListrangenum_hidden_layerslayergradient_checkpointingr7r9rr<r$s   zBlipTextEncoder.__init__NFTrvrwrxryrzrq use_cacher{output_hidden_states return_dictr|r@c  Csr|jr|jr|rtdd}|rAt|tr!tdt|}n t|t r0t|t |j d}n|durAtt |j dt |j d}| rEdnd} |rKdnd} |rU|durUdnd}t |j j D];}|j |}| rk| |f} |durs||nd}||||||||| }|d}|r| |df} |dur||df}q]| r| |f} | std d ||| | |fDSt||| | |d S) NzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...FzPassing 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)`.rr;rrrRcss|] }|dur|VqdSrcr;)rvr;r;r< sz*BlipTextEncoder.forward..)last_hidden_staterqrv attentionscross_attentions)rtrainingloggerwarningrr warning_oncer from_legacy_cacher r6rrrr)r8rvrwrxryrzrqrr{rrr|all_hidden_statesall_self_attentionsall_cross_attentionsr layer_modulelayer_head_mask layer_outputsr;r;r<rEsx           zBlipTextEncoder.forward) NNNNNNFFTN)rFrGrHr$r2rrrKr rrrrrErMr;r;r9r<rsJ       rcr)BlipTextPoolercs*tt|j|j|_t|_dSrc)r#r$rrYr'rTanh activationr7r9r;r<r$s zBlipTextPooler.__init__rvr@cCs(|dddf}||}||}|S)Nr)rr)r8rvfirst_token_tensor pooled_outputr;r;r<rEs  zBlipTextPooler.forwardrr;r;r9r<rs rcr)BlipTextPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dSr) r#r$rrYr'rrrrr transform_act_fnr,r-r7r9r;r<r$s  z(BlipTextPredictionHeadTransform.__init__rvr@cCs"||}||}||}|Src)rrr,rr;r;r<rEs   z'BlipTextPredictionHeadTransform.forwardrr;r;r9r<rs  rcs,eZdZfddZddZddZZS)BlipTextLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)bias) r#r$r transformrrYr'r&decoder Parameterr2zerosrr7r9r;r<r$s  z!BlipTextLMPredictionHead.__init__cCs|j|j_dSrc)rrrir;r;r< _tie_weights&sz%BlipTextLMPredictionHead._tie_weightscCrrc)rrrr;r;r<rE)rz BlipTextLMPredictionHead.forward)rFrGrHr$r rErMr;r;r9r<rs  rcr)BlipTextOnlyMLMHeadcstt||_dSrc)r#r$r predictionsr7r9r;r<r$1s zBlipTextOnlyMLMHead.__init__sequence_outputr@cCs||}|Src)r )r8r prediction_scoresr;r;r<rE5s zBlipTextOnlyMLMHead.forwardrr;r;r9r<r 0s r c@s*eZdZUdZeed<dZgZddZdS)BlipTextPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. r6bertcCs~t|tjtjfr|jjjd|jjdnt|tj r(|j j |jj dt|tjr;|j dur=|j j dSdSdS)zInitialize the weightsg)meanstd?N) rrrYr%weightdatanormal_r6initializer_ranger,rzero_fill_)r8moduler;r;r< _init_weightsEs  z%BlipTextPreTrainedModel._init_weightsN) rFrGrHrIr__annotations__base_model_prefix_no_split_modulesrr;r;r;r<r;s  rc#s"eZdZdZd#fdd ZddZddZd d Zd ed e e d e de def ddZ               d$deejd eejdeejdeejdeejdeejdeejdeejdeedee dee dee dee dee d eejdee ejeff d!d"ZZS)% BlipTextModela& The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of cross-attention is added between the self-attention layers, following the architecture described in [Attention is all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. argument and `is_decoder` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. TcsDt|||_t||_t||_|rt|nd|_| dSrc) r#r$r6rrDrencoderrpooler post_init)r8r6add_pooling_layerr9r;r<r$\s    zBlipTextModel.__init__cCs|jjSrcrDr)rir;r;r<get_input_embeddingsfsz"BlipTextModel.get_input_embeddingscCs ||j_dSrcr#)r8r]r;r;r<set_input_embeddingsis z"BlipTextModel.set_input_embeddingscCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsrrrr)r8heads_to_prunerrr;r;r< _prune_headsmszBlipTextModel._prune_headsrwrBrrr@c Cs^|dkr|dddddddf}n|dkr|r|\}}tj||d}|ddddf||d|ddddfk} | |j} | jd|jdkrl|jd| jd} tjtj||| f|| jd| gdd} | dddddddf|ddddddf}n|ddddddf}n t d |d |jd |j|jd }d |d}|S)a= Makes broadcastable attention and causal masks so that future and masked tokens are ignored. Arguments: attention_mask (`torch.Tensor`): Mask with ones indicating tokens to attend to, zeros for tokens to ignore. input_shape (`tuple[int]`): The shape of the input to the model. device (`torch.device`): The device of the input to the model. Returns: `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. rNrRrr)rr~r)axisz!Wrong shape for input_ids (shape z) or attention_mask (shape )rrg) rr2r3repeatrr~rcatonesrU) r8rwrBrrextended_attention_maskrrCseq_ids causal_maskprefix_seq_lenr;r;r<get_extended_attention_maskus4  . 6  z)BlipTextModel.get_extended_attention_maskNFr=rrxr>encoder_embedsryrzrqrr{rrr|c s| dur| njj} | dur| njj} | dur| njj} |r+| dur&| njj} nd} |dur9|dur9td|durO|||}|\}}|j}n,|durc|dd}|\}}|j}n|durw|dd}|\}}|j}ntdd}| durt | t s| ddj dn| }|durt |||f|}||||}|durt |tr|d\}}}n|\}}}||f}t |tr؇fdd |D}n|durt j||d }|}n|}nd}|jj}|durj||||d }n|}j|||||| | | | | |d }|d}jdur)|nd}| s8||f|d dSt|||j|j|j|jdS)a encoder_hidden_states (`torch.FloatTensor`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder. encoder_attention_mask (`torch.FloatTensor`, *optional*): Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. past_key_values (`Cache`, *optional*): Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up 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)`. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). NFzDYou cannot specify both input_ids and inputs_embeds at the same timerzGYou have to specify either input_ids or inputs_embeds or encoder_embedsrr}csg|]}|qSr;)invert_attention_mask)rmaskrir;r<rrz)BlipTextModel.forward..r))r=rr>r?) rwrxryrzrqrr{rrr|r)r pooler_outputrqrvrr)r6r{ruse_return_dictrrU%warn_if_padding_and_no_attention_maskrArrr rget_seq_lengthr2r.rr3listr5 get_head_maskrrDrr rrqrvrr)r8r=rwrrxr>r4ryrzrqrr{rrrr|rBrrCrr?r/encoder_batch_sizeencoder_sequence_lengthrencoder_hidden_shapeencoder_extended_attention_maskembedding_outputencoder_outputsr rr;rir<rEs$       zBlipTextModel.forward)T)NNNNNNNNNNNNNFN)rFrGrHrIr$r$r%r(rrrLrrr3rr2r rrrErMr;r;r9r<rSs  @     rc's,eZdZddgZfddZddZddZd d Zd d Z    d'de e j de e j de e j de e j de e j de e j de e j de e j de e de e de e de e de e de e de e d e ed!e e j d"eee j eff$d#d$Zd(fd%d& ZZS))BlipTextLMHeadModelzcls.predictions.decoder.weightzcls.predictions.decoder.biascs0t|t|dd|_t||_|j|_dS)NF)r")r#r$rrr clslabel_smoothingr7r9r;r<r$Js   zBlipTextLMHeadModel.__init__cCs |jSrc)rr$rir;r;r<r$Qrgz(BlipTextLMHeadModel.get_input_embeddingscCs|j|dSrc)rr%r8new_embeddingsr;r;r<r%Tsz(BlipTextLMHeadModel.set_input_embeddingscCs |jjjSrc)rDr rrir;r;r<get_output_embeddingsWrgz)BlipTextLMHeadModel.get_output_embeddingscCs||jj_|j|jj_dSrc)rDr rrrFr;r;r<set_output_embeddingsZs z)BlipTextLMHeadModel.set_output_embeddingsNFTrr=rwrrxr>ryrzlabelsrqrr{rr return_logitsr reductionr|r@cCsV| dur| n|jj} |durd} |j|||||||| | | | | ||d}|d}||}|r=|ddddddfSd}|dur|ddddddf}|ddddf|j}t||jd}|| d|jj | d}|dkr| | dd d}| s|f|d d}|dur|f|S|St |||j|j|j|jd S) a encoder_hidden_states (`torch.FloatTensor`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is configured as a decoder. encoder_attention_mask (`torch.FloatTensor`, *optional*): Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. labels (`torch.LongTensor`, *optional*): Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` past_key_values (`Cache`, *optional*): Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up 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)`. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). NF) rwrrxr>ryrzrqrr{rrrr|rrr)rLrEnonerR)losslogitsrqrvrr)r6r8rrDrrrrrErr&rAsumrrqrvrr)r8r=rwrrxr>ryrzrJrqrr{rrrKrrLr|rr r lm_lossshifted_prediction_scoresloss_fctrr;r;r<rE^sT*  zBlipTextLMHeadModel.forwardc s&tj|f||d|}d|d<|S)N)rqrwTr)r#prepare_inputs_for_generation)r8r=rqrw model_kwargs model_inputsr9r;r<rTsz1BlipTextLMHeadModel.prepare_inputs_for_generation)NNNNNNNNNNNNNFTrN)NN)rFrGrH_tied_weights_keysr$r$r%rHrIrr2rr rrrrrrErTrMr;r;r9r<rCGs|       \rC)rrCr)8rtypingrrr2rrrtorch.nnr activationsr cache_utilsr r r generationr modeling_layersrmodeling_outputsrrrmodeling_utilsr pytorch_utilsrrrutilsrutils.deprecationrconfiguration_blipr get_loggerrFrModulerrNrrrrrrrrrr rrrC__all__r;r;r;r<sF         43>b u