o i@sdZddlZddlmZddlmZmZmZmZddl Z ddl m Z ddl m Z ddlmZddlmZmZmZmZdd lmZmZdd lmZmZmZdd lmZmZmZm Z m!Z!m"Z"d d l#m$Z$m%Z%m&Z&e!'e(Z)de j*de j*fddZ+de j*de j*fddZ,eeGdddeZ-Gddde j.Z/Gddde j.Z0Gddde j.Z1de0iZ2Gddde j.Z3Gd d!d!e j.Z4Gd"d#d#e j.Z5Gd$d%d%eZ6Gd&d'd'e j.Z7Gd(d)d)e j.Z8 *dOd+e j.d,e j*d-e j*d.e j*d/ee j*d0e9d1e9fd2d3Z:Gd4d5d5e j.Z;Gd6d7d7e j.ZGdd?d?eZ@Gd@dAdAe j.ZAGdBdCdCe@ZBedDdEGdFdGdGe@ZCGdHdIdIe@ZDGdJdKdKe@ZEdPdLdMZFgdNZGdS)QzPyTorch AltCLIP model.N) dataclass)AnyCallableOptionalUnion)ACT2FN)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPooling,BaseModelOutputWithPoolingAndCrossAttentions'BaseModelOutputWithPoolingAndProjection)ALL_ATTENTION_FUNCTIONSPreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputauto_docstringcan_return_tuplefilter_out_non_signature_kwargslogging torch_int) AltCLIPConfigAltCLIPTextConfigAltCLIPVisionConfiglogitsreturncCstj|tjt||jdS)Ndevice)nn functional cross_entropytorcharangelenr )rr'i/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/transformers/models/altclip/modeling_altclip.pycontrastive_loss+sr) similaritycCs t|}t|}||dS)Ng@)r)t)r* caption_loss image_lossr'r'r( clip_loss/s  r.c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<dZ eed<dZeed <d eefd d ZdS) AltCLIPOutputa loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Contrastive loss for image-text similarity. logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text similarity scores. logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image similarity scores. text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPTextModel`]. image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPVisionModel`]. text_model_output (`BaseModelOutputWithPooling`): The output of the [`AltCLIPTextModel`]. vision_model_output (`BaseModelOutputWithPooling`): The output of the [`AltCLIPVisionModel`]. Nlosslogits_per_imagelogits_per_text text_embeds image_embedstext_model_outputvision_model_outputrcstfddDS)Nc3s.|]}|dvr |nt|VqdS))r5r6N)getattrto_tuple).0kselfr'r( Us  z)AltCLIPOutput.to_tuple..)tuplekeysr;r'r;r(r8Ts zAltCLIPOutput.to_tuple)__name__ __module__ __qualname____doc__r0rr$ FloatTensor__annotations__r1r2r3r4r5r r6r>rr8r'r'r'r(r/5s   r/cs4eZdZdZfddZ d ddZdd ZZS) AltRobertaEmbeddingszV Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_t|dd|_|jdt|jddd|jd tj|jtjd dd|j|_tj|j|j|jd|_ dS) N) padding_idxepsposition_embedding_typeabsolute position_idsrF persistenttoken_type_idsdtype)super__init__r! Embedding vocab_size hidden_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutr7rJregister_bufferr$r%expandzerosrLsizelongrGr<config __class__r'r(rUbs"   zAltRobertaEmbeddings.__init__Nrc Cs|dur|durt||j|}n||}|dur|}n|dd}|d}|durTt|drI|jddd|f}||d|} | }n tj|tj |j j d}|dur]| |}| |} || } |jdkrt||} | | 7} || } || } | S)NrNrrQrrSr rK)"create_position_ids_from_input_idsrG&create_position_ids_from_inputs_embedsrghasattrrQrer$rfrhrLr rZr^rJr\r_rc) r< input_idsrQrL inputs_embedspast_key_values_length input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr^ embeddingsr\r'r'r(forward{s0         zAltRobertaEmbeddings.forwardcCsN|dd}|d}tj|jd||jdtj|jd}|d|S)z We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids. Args: inputs_embeds: torch.Tensor Returns: torch.Tensor NrNrrmr)rgr$r%rGrhr unsqueezere)r<rrrtsequence_lengthrLr'r'r(ros  z;AltRobertaEmbeddings.create_position_ids_from_inputs_embeds)NNNNr)r@rArBrCrUryro __classcell__r'r'rkr(rF\s   (rFc s\eZdZd fdd Z   d dejdeejdeejdeed e ejf d d Z Z S)AltRobertaSelfAttentionNcst|j|jdkrt|dstd|jd|jd|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_|p\t|dd|_|jdksh|jd kr{|j|_t d |jd |j|_dSdS) Nrembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()rJrK relative_keyrelative_key_queryr)rTrUrXnum_attention_headsrp ValueErrorintattention_head_size all_head_sizer!Linearquerykeyvalueraattention_probs_dropout_probrcr7rJr[rVdistance_embeddingr<rjrJrkr'r(rUs*  z AltRobertaSelfAttention.__init__F hidden_statesattention_mask head_maskoutput_attentionsrcCs|jdd}g|d|jR}|||dd}|||dd}|||dd} t||dd} |j dksI|j dkr|jd|jd} } tj | tj |j ddd} tj | tj |j ddd}| |}| ||jd}|j|jd}|j dkrtd ||}| |} n|j dkrtd ||}td ||}| ||} | t|j} |dur| |} tjj| dd }||}|dur||}t|| }|d ddd }|dd|jf}||}|r||f}|S|f}|S)NrNrrrrrmrRzbhld,lrd->bhlrzbhrd,lrd->bhlrdimrr)shaperrview transposerrr$matmulrJr%rhr rr[torSeinsummathsqrtr!r"softmaxrcpermute contiguousrgr)r<rrrrrt hidden_shape query_layer key_layer value_layerattention_scores query_length key_lengthposition_ids_lposition_ids_rdistancepositional_embeddingrelative_position_scoresrelative_position_scores_queryrelative_position_scores_keyattention_probs context_layernew_context_layer_shapeoutputsr'r'r(rysF        zAltRobertaSelfAttention.forwardNNNF) r@rArBrUr$TensorrrDboolr>ryr|r'r'rkr(r}s r}c8eZdZfddZdejdejdejfddZZS)AltRobertaSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrH) rTrUr!rrXdenser_r`rarbrcrirkr'r(rU  zAltRobertaSelfOutput.__init__r input_tensorrcC&||}||}|||}|Srrrcr_r<rrr'r'r(ry  zAltRobertaSelfOutput.forwardr@rArBrUr$rryr|r'r'rkr(r  $reagerc sdeZdZdfdd ZddZ   ddejdeejd eejd ee d e ejf d d Z Z S)AltRobertaAttentionNcs4tt|j||d|_t||_t|_dS)N)rJ) rTrU"ALT_ROBERTA_SELF_ATTENTION_CLASSES_attn_implementationr<routputset pruned_headsrrkr'r(rU s   zAltRobertaAttention.__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)r&rr<rrrrrrrrrrunion)r<headsindexr'r'r( prune_heads(s zAltRobertaAttention.prune_headsFrrrrrcCs8|j||||d}||d|}|f|dd}|SN)rrrrr)r<r)r<rrrr self_outputsattention_outputrr'r'r(ry:szAltRobertaAttention.forwardrr) r@rArBrUrr$rrrDrr>ryr|r'r'rkr(rs"rc2eZdZfddZdejdejfddZZS)AltRobertaIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr) rTrUr!rrXintermediate_sizer isinstance hidden_actstrrintermediate_act_fnrirkr'r(rUNs   zAltRobertaIntermediate.__init__rrcCs||}||}|Sr)rrr<rr'r'r(ryVs  zAltRobertaIntermediate.forwardrr'r'rkr(rMs rcr)AltRobertaOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSr) rTrUr!rrrXrr_r`rarbrcrirkr'r(rU^rzAltRobertaOutput.__init__rrrcCrrrrr'r'r(rydrzAltRobertaOutput.forwardrr'r'rkr(r]rrc sbeZdZfddZ   ddejdeejdeejdeed e ejf d d Z d d Z Z S)AltRobertaLayercs:t|j|_d|_t||_t||_t||_ dS)Nr) rTrUchunk_size_feed_forward seq_len_dimr attentionr intermediaterrrirkr'r(rUms   zAltRobertaLayer.__init__NFrrrrrc KsP|j|f|||d|}|d}|dd}t|j|j|j|} | f|}|Sr)rrfeed_forward_chunkrr) r<rrrrkwargsself_attention_outputsrr layer_outputr'r'r(ryus   zAltRobertaLayer.forwardcCs||}|||}|Sr)rr)r<rintermediate_outputrr'r'r(rs  z"AltRobertaLayer.feed_forward_chunkr) r@rArBrUr$rrrDrr>ryrr|r'r'rkr(rls"   rcszeZdZfddZe     ddejdeejdeejd ee d ee d ee d e e eje ffd dZ ZS)AltRobertaEncoderc:t|_tfddtjD|_d|_dS)Ncg|]}tqSr')r)r9irjr'r( z.AltRobertaEncoder.__init__..F) rTrUrjr! ModuleListrangenum_hidden_layerslayergradient_checkpointingrirkrr(rU   zAltRobertaEncoder.__init__NFTrrrroutput_hidden_states return_dictrcKs|rdnd}|r dnd} t|jD].\} } |r||f}|dur$|| nd} | d||| |d|} | d}|r?| | df} q|rG||f}t||| dS)Nr')rrrrrrlast_hidden_stater attentions) enumeraterr )r<rrrrrrrall_hidden_statesall_self_attentionsr layer_modulelayer_head_mask layer_outputsr'r'r(rys2   zAltRobertaEncoder.forward)NNFFT)r@rArBrUrr$rrrDrrr>r ryr|r'r'rkr(rs.  rcr)AltRobertaPoolercs*tt|j|j|_t|_dSr)rTrUr!rrXrTanh activationrirkr'r(rUs zAltRobertaPooler.__init__rrcCs(|dddf}||}||}|S)Nr)rr)r<rfirst_token_tensor pooled_outputr'r'r(rys  zAltRobertaPooler.forwardrr'r'rkr(rs rmodulerrrrscalingrcc Ks|t||dd|}|dur||}tjj|dtjd|j}tjj |||j d}t||} | dd } | |fS)NrNr)rrS)ptrainingrr) r$rrr!r"rfloat32rrSrcrr) rrrrrrrcr attn_weights attn_outputr'r'r(eager_attention_forwards  r csheZdZdZfddZ   d dejdeejdeejd eed e ejeejff d d Z Z S)AltCLIPAttentionz=Multi-headed attention from 'Attention Is All You Need' papercst||_|j|_|j|_|j|j|_|j|j|jkr-td|jd|jd|jd|_ |j |_ d|_ t |j|j|_t |j|j|_t |j|j|_t |j|j|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: ).F)rTrUrjrX embed_dimr num_headshead_dimrscaleattention_dropoutrc is_causalr!rk_projv_projq_projout_projrirkr'r(rUs$   zAltCLIPAttention.__init__NFrrcausal_attention_maskrrc Cs0|j\}}}||}||} ||} ||||j|jdd}| |||j|jdd} | |||j|jdd} |jj dkrY|durR|durR||}n |durX|}n|du|_ t } |jj dkrlt |jj } | ||| | ||j |j |js{dn|jd\} } | |||} || } |sd} | | fS)z#Input shape: Batch x Time x Channelrrflash_attention_2Nrr)rrrc)rrrrrrrrrjrrr rrrrcreshaperr)r<rrrr batch_sizerur queriesr?valuesattention_interfacerrr'r'r(rys@          zAltCLIPAttention.forwardr) r@rArBrCrUr$rrrr>ryr|r'r'rkr(r s" r cr) AltCLIPMLPcsDt||_t|j|_t|j|j |_ t|j |j|_ dSr) rTrUrjrr activation_fnr!rrXrfc1fc2rirkr'r(rU9s  zAltCLIPMLP.__init__rrcCs"||}||}||}|Sr)r rr!rr'r'r(ry@s   zAltCLIPMLP.forwardrr'r'rkr(r8s rc sTeZdZdeffdd Z d dejdejdejdeed e ej f d d Z Z S) AltCLIPEncoderLayerrjcsRt|j|_t||_tj|j|jd|_ t ||_ tj|j|jd|_ dSr) rTrUrXr r  self_attnr!r_r` layer_norm1rmlp layer_norm2rirkr'r(rUHs   zAltCLIPEncoderLayer.__init__FrrrrrcCsd|}||}|j||||d\}}||}|}||}||}||}|f}|r0||f7}|S)aI 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. `(config.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. )rrrr)r$r#r&r%)r<rrrrresidualrrr'r'r(ryPs"     zAltCLIPEncoderLayer.forwardF) r@rArBrrUr$rrrr>rDryr|r'r'rkr(r"Gs r"csxeZdZdZdeffdd Ze     ddeej deej dee d ee d ee d e e e ff d d ZZS)AltCLIPEncoderz Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`AltCLIPEncoderLayer`]. Args: config: AltCLIPConfig rjcr)Ncrr')r")r9_rr'r(rrz+AltCLIPEncoder.__init__..F) rTrUrjr!rrrlayersrrirkrr(rUrzAltCLIPEncoder.__init__Nrrrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|r"dnd}|r(dnd}|} t|jD] \} } |r<|| f}| | |||d} | d} |rQ|| df}q1|rY|| f}t| ||dS)a Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): 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. 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) causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Causal mask for the text model. 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) 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. Nr')rrrr)rjrruse_return_dictrr+r ) r<rrrrrrrencoder_statesall_attentionsridx encoder_layerrr'r'r(rys2'    zAltCLIPEncoder.forward)NNNNN)r@rArBrCrrUrrr$rrrr>r ryr|r'r'rkr(r)ys, r)csXeZdZdeffdd Zdejdededejfdd Zdd ej dejfd d Z Z S)AltCLIPVisionEmbeddingsrjcst||_|j|_|j|_|j|_tt |j|_ tj |j |j|j|jdd|_|j|jd|_|jd|_t|j|j|_|jdt |jddddS)NF) in_channels out_channels kernel_sizestridebiasrrrLrMrO)rTrUrjrXr  image_size patch_sizer! Parameterr$randnclass_embeddingConv2d num_channelspatch_embedding num_patches num_positionsrVposition_embeddingrdr%rerirkr'r(rUs"  "z AltCLIPVisionEmbeddings.__init__rxheightwidthrc Cs|jdd}|jjd}|jdd}tjs(||kr(||kr(||jS|ddddf}|ddddf}|jd} ||j} ||j} t |d} | d| | | }| dddd}t j j|| | fdd d }| dddddd| }tj||fdd S) a This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This method is also adapted to support torch.jit tracing. Adapted from: - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 rrNrNg?rrbicubicF)rgmode align_cornersr)rrAweightrzr$jit is_tracingrLr8rrrr!r" interpolatercat) r<rxrBrCr?rAr@class_pos_embedpatch_pos_embedr new_height new_widthsqrt_num_positionsr'r'r(interpolate_pos_encodings*      z0AltCLIPVisionEmbeddings.interpolate_pos_encodingF pixel_valuesc Cs|j\}}}}|s&||jks||jkr&td|d|d|jd|jd |jjj}||j|d}|ddd}|j |dd} t j | |gdd } |r[| | | ||} | S| ||j} | S) NzInput image size (*z) doesn't match model (r rRrrrNr)rr7rr>rGrSrflattenrr;rer$rKrQrArL) r<rRrQrr*rBrC target_dtype patch_embeds class_embedsrxr'r'r(rys  zAltCLIPVisionEmbeddings.forwardr() r@rArBrrUr$rrrQrDryr|r'r'rkr(r1s )r1c@s*eZdZUeed<dZdZgZddZdS)AltCLIPPreTrainedModelrjaltclipTcCs|jj}t|tr:|jj}tjj|jd|jd|dtjj|j j |jj |dtjj|j j |jj |ddSt|t r|jj}|jdd|jjd|}|jd|}tjj|jj |dtjj|jj |dtjj|jj |dtjj|jj |ddSt|tr|jj}|jjdd|jjd|}d|jjd|}tjj|jj |dtjj|jj |ddSt|trtjj|jj |jd|jjdd|j_tjj|jj |jd|jjdd|j_dSt|tjr|jj |j j!ddSt|tj"r%|j jjd|jjd|jdur#|jj dSdSt|tj#rH|j jjd|jjd|j$durJ|j j|j$ dSdSdS) zInitialize the weightsrr )meanstd)r[rTg?N)%rjinitializer_factorrr1r!initnormal_r;r r>rGinitializer_rangerAr rrrrrrrXr r! AltCLIPModeltext_projectiontext_embed_dim_is_hf_initializedvisual_projectionvision_embed_dimr_r6datazero_fill_rrVrG)r<rfactor in_proj_std out_proj_stdfc_stdr'r'r( _init_weights+sZ          z$AltCLIPPreTrainedModel._init_weightsN) r@rArBrrEbase_model_prefixsupports_gradient_checkpointing_no_split_modulermr'r'r'r(rX$s  rXcsveZdZdeffdd Zee     ddeej dee dee d ee d ee d e e e ff d d ZZS)AltCLIPVisionTransformerrjcsRt||_|j}t||_tj||jd|_ t ||_ tj||jd|_ dSr) rTrUrjrXr1rxr!r_r` pre_layrnormr)encoderpost_layernorm)r<rjr rkr'r(rUZs   z!AltCLIPVisionTransformer.__init__NFrRrrrrQrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|j||d}||}|j|||dd}|d}|dddddf} || } t || |j |j dS)Nz You have to specify pixel_values)rQT)rrrrrrr pooler_outputrr) rjrrr,rrxrrrsrtr rr) r<rRrrrrQrencoder_outputsrrr'r'r(ryds.   z AltCLIPVisionTransformer.forward)NNNNF)r@rArBrrUrrrr$rDrrr>r ryr|r'r'rkr(rqYs,  rqcseZdZUeed<dZdeffdd ZdejfddZ e    dde e j d e ed e ed ed e edeeeff ddZZS)AltCLIPVisionModelrjrRcs"t|t||_|dSr)rTrUrq vision_model post_initrirkr'r(rUs   zAltCLIPVisionModel.__init__rcC |jjjSr)ryrxr>r;r'r'r(get_input_embeddings z'AltCLIPVisionModel.get_input_embeddingsNFrrrQrcCs(|dur|n|jj}|j|||||dS)a Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, AltCLIPVisionModel >>> model = AltCLIPVisionModel.from_pretrained("BAAI/AltCLIP") >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooled_output = outputs.pooler_output # pooled CLS states ```NrRrrrQr)rjr,ry)r<rRrrrQrr'r'r(ryszAltCLIPVisionModel.forward)NNNFN)r@rArBrrEmain_input_namerUr!Moduler|rrr$rDrrr>r ryr|r'r'rkr(rxs0  rxaE The model behaves as an encoder following the architecture described in *Attention is all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. .. _*Attention is all you need*: https://huggingface.co/papers/1706.03762 ) custom_introcseZdZUeed<dfdd ZddZddZd d Ze dd e e j d e e j de e j de e j de e j de e j de e de e de e deee j effddZZS)AltRobertaModelrjTcsDt|||_t||_t||_|rt|nd|_| dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) rTrUrjrFrxrrsrpoolerrz)r<rjadd_pooling_layerrkr'r(rUs    zAltRobertaModel.__init__cCs|jjSrrxrZr;r'r'r(r|sz$AltRobertaModel.get_input_embeddingscCs ||j_dSrrr<rr'r'r(set_input_embeddings z$AltRobertaModel.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)itemsrsrrr)r<heads_to_prunerrr'r'r( _prune_headsszAltRobertaModel._prune_headsNrqrrQrLrrrrrrrc Cs|dur|n|jj}|dur|n|jj}| dur| n|jj} |dur*|dur*td|dur9||||} n|durF|dd} ntd| \} } |durU|jn|j} |duretj | | f| d}|durt |j dr|j j ddd| f}| | | }|}n tj| tj| d}||| }|||jj}|j ||||d}|j|||||dd }|d }|jdur||nd}t|||j|jd S) NzDYou cannot specify both input_ids and inputs_embeds at the same timerNz5You have to specify either input_ids or inputs_embedsrrQrm)rqrLrQrrT)rrrrrrru)rjrrr,r%warn_if_padding_and_no_attention_maskrgr r$onesrprxrQrerfrhget_extended_attention_mask get_head_maskrrsrr rr)r<rqrrQrLrrrrrrrtrrur rvrwextended_attention_maskembedding_outputrwsequence_outputrr'r'r(rys\     zAltRobertaModel.forward)T NNNNNNNNN)r@rArBrrErUr|rrrrr$rrrr>r ryr|r'r'rkr(rsJ     rcseZdZUeed<fddZdejfddZdej ddfd d Z dd e e dej ffd d Z ee         dde ejde ejde ejde ejde ejde ejde ede ede edeeeffddZZS)AltCLIPTextModelrjcsLt|t|dd|_t|j|j|_tj |j|j d|_ | dS)NF)rrH) rTrUrrobertar!rrX project_dimtransformationr_r`pre_LNrzrirkr'r(rU<s  zAltCLIPTextModel.__init__rcCr{rrrxrZr;r'r'r(r|Cr}z%AltCLIPTextModel.get_input_embeddingsrNcCs||jj_dSrrrr'r'r(rFsz%AltCLIPTextModel.set_input_embeddingsnew_num_tokenscs t|Sr)rTresize_token_embeddings)r<rrkr'r(rIrz(AltCLIPTextModel.resize_token_embeddingsrqrrQrLrrrrrrc  Csp|dur|n|jj}|j|||||||| dd } | d} || } || } | dddf} t| | | j| jdS)a+ Examples: ```python >>> from transformers import AutoProcessor, AltCLIPTextModel >>> model = AltCLIPTextModel.from_pretrained("BAAI/AltCLIP") >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP") >>> texts = ["it's a cat", "it's a dog"] >>> inputs = processor(text=texts, padding=True, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_state = outputs.last_hidden_state >>> pooled_output = outputs.pooler_output # pooled CLS states ```NT) rqrrQrLrrrrrrrru)rjr,rrrr rr)r<rqrrQrLrrrrrrrrprojection_statervr'r'r(ryLs,   zAltCLIPTextModel.forwardrr)r@rArBrrErUr!rr|rVrrrrrrr$rrrr>r ryr|r'r'rkr(r9sL     rcs eZdZUeed<deffdd Zee   ddej de ej de ej de ej d ej f d d Z ee dd ej de d ej fddZe         dde ejd e ej de ej de ejde ej de e de e de e de de e d eeeffddZZS)r`rjcst|t|jtstdt|jdt|jts(tdt|jd|j}|j}|j |_ |j |_ |j |_ |j |_t||_t||_tj|j|j dd|_tj|j |j dd|_tt|jj|_|dS)NzRconfig.vision_config is expected to be of type AltCLIPVisionConfig but is of type .zNconfig.text_config is expected to be of type AltCLIPTextConfig but is of type F)r6)rTrUr vision_configr TypeErrortype text_configrrprojection_dimrrbrXrer text_modelrqryr!rrdrar9r$tensorrjlogit_scale_init_value logit_scalerz)r<rjrrrkr'r(rUs2      zAltCLIPModel.__init__NrqrrLrQrcCs&|j||||d}|j}||}|S)a Returns: text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPTextModel`]. Examples: ```python >>> import torch >>> from transformers import AutoProcessor, AltCLIPModel >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP") >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP") >>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") >>> with torch.inference_mode(): ... text_features = model.get_text_features(**inputs) ```)rqrrLrQ)rrvra)r<rqrrLrQ text_outputsr text_featuresr'r'r(get_text_featuress zAltCLIPModel.get_text_featuresFrRrQcCs"|j||d}|j}||}|S)aQ Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPVisionModel`]. Examples: ```python >>> import torch >>> from transformers import AutoProcessor, AltCLIPModel >>> from transformers.image_utils import load_image >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP") >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = load_image(url) >>> inputs = processor(images=image, return_tensors="pt") >>> with torch.inference_mode(): ... image_features = model.get_image_features(**inputs) ```)rRrQ)ryrvrd)r<rRrQvision_outputsrimage_featuresr'r'r(get_image_featuress zAltCLIPModel.get_image_features return_lossrrrc  Cs(|dur|n|jj}|dur|n|jj}| dur| n|jj} |j||||||| d} |j|||| | d} | d} || } | d}||}| | jdddd} ||jdddd}|j }t || |}|j}d}|rtt|}| s|||| | | f}|dur|f|S|St||||| | | d S) a return_loss (`bool`, *optional*): Whether or not to return the contrastive loss. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, AltCLIPModel >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP") >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor( ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True ... ) >>> outputs = model(**inputs) >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities ```N)rqrrQrLrrrr~rrrNT)rrkeepdim)r0r1r2r3r4r5r6)rjrrr,rryrdranormrexpr$rr+Tr.r/)r<rqrRrrLrQrrrrQrrrr4r3rr2r1r0rr'r'r(rysX%    zAltCLIPModel.forward)NNNr() NNNNNNNNFN)r@rArBrrErUrrr$rrrDrrr LongTensorrr>r/ryr|r'r'rkr(r`s~ !$$    r`cCs6||}tj|dd|||}||S)a Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols are ignored. This is modified from fairseq's `utils.make_positions`. Args: x: torch.Tensor x: Returns: torch.Tensor rr)nerr$cumsumtype_asrh)rqrGrsmaskincremental_indicesr'r'r(rn\s  rn)rXrxrr`)r)r)HrCr dataclassesrtypingrrrrr$torch.nnr! activationsrmodeling_layersr modeling_outputsr r r r modeling_utilsrr pytorch_utilsrrrutilsrrrrrrconfiguration_altcliprrr get_loggerr@loggerrr)r.r/rrFr}rrrrrrrrfloatr r rr"r)r1rXrqrxrrr`rn__all__r'r'r'r(s      $YW.)2 J2XS445 nS Q