U
    9%eY                    @   s&  d Z ddlZddlmZ ddlmZ ddlmZmZm	Z	m
Z
 ddlZddlZddlmZ ddlmZ dd	lmZmZ dd
lmZmZmZmZmZ ddlmZmZ ddlmZmZ ddlm Z m!Z!m"Z"m#Z# ddl$m%Z%m&Z&m'Z' e"(e)Z*dZ+dZ,dZ-ddgZ.dZ/dZ0eG dd deZ1eG dd deZ2G dd dej3Z4G dd dej3Z5G dd dej3Z6G d d! d!ej3Z7G d"d# d#ej3Z8G d$d% d%ej3Z9G d&d' d'ej3Z:G d(d) d)ej3Z;G d*d+ d+ej3Z<G d,d- d-ej3Z=G d.d/ d/ej3Z>G d0d1 d1ej3Z?G d2d3 d3ej3Z@G d4d5 d5ej3ZAG d6d7 d7ej3ZBdTd8d9ZCG d:d; d;eZDG d<d= d=eDZEG d>d? d?eDZFe d@e/G dAdB dBeDZGG dCdD dDej3ZHG dEdF dFej3ZIG dGdH dHej3ZJe dIe/G dJdK dKeDZKe dLe/G dMdN dNeDZLG dOdP dPej3ZMe dQe/G dRdS dSeDZNdS )UzPyTorch BridgeTower Model    N)OrderedDict)	dataclass)ListOptionalTupleUnion)nn)CrossEntropyLoss   )ACT2FNQuickGELUActivation))BaseModelOutputWithPastAndCrossAttentions,BaseModelOutputWithPoolingAndCrossAttentionsMaskedLMOutputModelOutputSequenceClassifierOutput)PreTrainedModelapply_chunking_to_forward) find_pruneable_heads_and_indicesprune_linear_layer)add_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings   )BridgeTowerConfigBridgeTowerTextConfigBridgeTowerVisionConfigr   zBridgeTower/bridgetower-baseZRobertaTokenizerz$BridgeTower/bridgetower-base-itm-mlmaO  
    This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ subclass. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`BridgeTowerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
a<  
    Args:
        input_ids (`torch.LongTensor` of shape `({0})`):
            Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
            [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
            IDs?](../glossary#input-ids)

        attention_mask (`torch.FloatTensor` of shape `({0})`, *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)

        token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
            Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
            1]`:
            - 0 corresponds to a *sentence A* token,
            - 1 corresponds to a *sentence B* token.
            [What are token type IDs?](../glossary#token-type-ids)

        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`BridgeTowerImageProcessor`]. See
            [`BridgeTowerImageProcessor.__call__`] for details.

        pixel_mask (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
            Mask to avoid performing attention on padding pixel values. Mask values selected in `[0, 1]`:

            - 1 for pixels that are real (i.e. **not masked**),
            - 0 for pixels that are padding (i.e. **masked**).
            `What are attention masks? <../glossary.html#attention-mask>`__

        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
            Mask to nullify selected heads of the self-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 `({0}, 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.

        image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
            Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `pixel_values` into patch embeddings.

        image_token_type_idx (`int`, *optional*):
            - The token type ids for images.

        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.
c                   @   sh   e Zd ZU dZdZejed< dZejed< dZ	ejed< dZ
eeej  ed< dZeeej  ed< dS )BridgeTowerModelOutputa  
    Output type of [`BridgeTowerModel`].

    Args:
        text_features (`torch.FloatTensor` of shape `(batch_size, text_sequence_length, hidden_size)`):
            Sequence of hidden-states at the text output of the last layer of the model.
        image_features (`torch.FloatTensor` of shape `(batch_size, image_sequence_length, hidden_size)`):
            Sequence of hidden-states at the image output of the last layer of the model.
        pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size x 2)`):
            Concatenation of last layer hidden-state of the first token of the text and image sequence (classification
            token), respectively, after further processing through layers used for auxiliary pretraining tasks.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
            the model at the output of each layer plus the optional initial embedding outputs.
        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
            sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.
    Ntext_featuresimage_featurespooler_outputhidden_states
attentions)__name__
__module____qualname____doc__r   torchFloatTensor__annotations__r    r!   r"   r   r   r#    r+   r+   s/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/bridgetower/modeling_bridgetower.pyr   }   s   
r   c                   @   s   e Zd ZU dZdZeej ed< dZ	ejed< dZ
eeej  ed< dZeeej  ed< dZeeej  ed< dZeeej  ed< dZeeej  ed	< dS )
BridgeTowerContrastiveOutputaZ  
    Output type of ['BridgeTowerForContrastiveLearning']

    Args:
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`:
            Image-text contrastive loss.
        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        text_embeds (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
            The text embeddings obtained by applying the projection layer to the pooler_output.
        image_embeds (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
            The image embeddings obtained by applying the projection layer to the pooler_output.
        cross_embeds  (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
            The text-image cross-modal embeddings obtained by applying the projection layer to the pooler_output.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of
            the model at the output of each layer plus the optional initial embedding outputs.
        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
            sequence_length)`.
    Nlosslogitstext_embedsimage_embedscross_embedsr"   r#   )r$   r%   r&   r'   r.   r   r(   r)   r*   r/   r0   r   r1   r2   r"   r#   r+   r+   r+   r,   r-      s   
r-   c                       sF   e Zd Z fddZejejdddZd	ejejdddZ  ZS )
BridgeTowerResidualAttentionc              
      s   t    t|j|jd | _tj|j|jd| _t	t
dt|j|jd fdt fdt|jd |jfg| _tj|j|jd| _d | _d S )N@   Zepsc_fc   Zgeluc_proj)super__init__r   ZMultiheadAttentionhidden_sizeattn	LayerNormlayer_norm_epsln_1Z
ModuleDictr   Linearr   mlpln_2	attn_maskselfconfig	__class__r+   r,   r:      s    
	z%BridgeTowerResidualAttention.__init__hidden_stateattention_maskc                 C   sZ   |d k	r|j tj|jd}| jd k	r8| jj |j|jdnd | _| j|||d| j|dd S )NdtypedeviceF)Zneed_weightsrC   Zkey_padding_maskr   )tor(   boolrN   rC   rM   r<   )rE   rJ   rK   r+   r+   r,   	attention   s     z&BridgeTowerResidualAttention.attentionNc                 C   sH   ||  | || }| |}| j D ]\}}||}q*|| }|S N)rQ   r?   rB   rA   items)rE   rJ   rK   Zresidual_state_layerr+   r+   r,   forward   s    

z$BridgeTowerResidualAttention.forward)N)	r$   r%   r&   r:   r(   TensorrQ   rV   __classcell__r+   r+   rG   r,   r3      s   r3   c                       s6   e Zd Z fddZdejeej dddZ  ZS )BridgeTowerTransformerc                    sr   t     j| _ j| _ jrFt fddt| jd D | _n t fddt| jD | _ j	| _	d S )Nc                    s   g | ]}t  qS r+   r3   .0rT   rF   r+   r,   
<listcomp>   s     z3BridgeTowerTransformer.__init__.<locals>.<listcomp>r   c                    s   g | ]}t  qS r+   rZ   r[   r]   r+   r,   r^      s     )
r9   r:   r;   num_hidden_layersZremove_last_layerr   
ModuleListrange	resblocksstop_gradientrD   rG   r]   r,   r:      s    
zBridgeTowerTransformer.__init__NrI   c                 C   s>   g }| j D ].}|||}| jr.||  q
|| q
|S rR   )rb   rc   appenddetach)rE   rJ   rK   r"   blockr+   r+   r,   rV      s    

zBridgeTowerTransformer.forward)N)	r$   r%   r&   r:   r(   rW   r   rV   rX   r+   r+   rG   r,   rY      s   rY   c                       s6   e Zd Zed fddZejejdddZ  Z	S )BridgeTowerVisionEmbeddingsr]   c                    s   t    || _|j| _|j| _|j| _tt	
| j| _tj|j| j| j| jdd| _| j| j d | _| jd | _t| j| j| _| jdt	| jddd d S )NF)Zin_channelsZout_channelsZkernel_sizeZstridebias   r   position_idsr   
persistent)r9   r:   rF   r;   Z	embed_dimZ
image_sizeZ
patch_sizer   	Parameterr(   Zrandnclass_embeddingConv2dZnum_channelspatch_embeddingZnum_patchesZnum_positions	Embeddingposition_embeddingregister_bufferarangeexpandrD   rG   r+   r,   r:     s"    
z$BridgeTowerVisionEmbeddings.__init__)pixel_valuesreturnc                 C   sn   |j d }| jjj}| |j|d}|ddd}| j|dd}t	j
||gdd}|| | j }|S )Nr   rM   ri   r   rl   dim)shaperr   weightrM   rO   flatten	transposerp   rw   r(   catrt   rj   )rE   rx   
batch_sizeZtarget_dtypeZpatch_embedsZclass_embeds
embeddingsr+   r+   r,   rV     s    

z#BridgeTowerVisionEmbeddings.forward)
r$   r%   r&   r   r:   r(   r)   rW   rV   rX   r+   r+   rG   r,   rg     s   rg   c                       sL   e Zd Z fddZejdddZejdddZejdd	d
Z  Z	S )BridgeTowerVisionTransformerc                    sx   t    t | _tj j jd| _t	 | _
tj j jd| _ j| _ jstt fddt jD | _d S )Nr5   c                    s   g | ]}t j j jd qS )r5   )r   r=   r;   r>   r[   r]   r+   r,   r^   4  s     z9BridgeTowerVisionTransformer.__init__.<locals>.<listcomp>)r9   r:   rg   r   r   r=   r;   r>   ln_prerY   transformerln_postshare_layernormr`   ra   r_   ln_separaterD   rG   r]   r,   r:   *  s    


z%BridgeTowerVisionTransformer.__init__)rx   c                 C   s   |  |}| |}|ddd}| ||}tj|dd}|dddd}| jr^| |}n:g }t|| j	D ]\}}||}|
| qntj|dd}|S )Nr   r   ri   r{   r
   )r   r   permuter   r(   stackr   r   zipr   rd   )rE   rx   rK   r"   Zhidden_states_stacklnr+   r+   r,   rV   7  s    

z$BridgeTowerVisionTransformer.forwardc                 C   s&   |  |}| |}|ddd}|S Nr   r   ri   )r   r   r   )rE   rx   r"   r+   r+   r,   forward_preM  s    

z(BridgeTowerVisionTransformer.forward_pre)rJ   c                 C   s   | ddd}| |}|S r   )r   r   )rE   rJ   Zvisual_output_postr+   r+   r,   forward_postT  s    
z)BridgeTowerVisionTransformer.forward_post)
r$   r%   r&   r:   r(   rW   rV   r   r   rX   r+   r+   rG   r,   r   )  s   r   c                       s$   e Zd Z fddZdd Z  ZS )BridgeTowerLinkTowerc                    s   t    |j| _|j| _|jdkrt|jdkrBttd| _n|jdkr^ttd| _	tj
| j|jd| _
ntd|j dd S )	N)add
scaled_addinterpolater         ?r   g      ?r5   link_tower_type  is not implemented)r9   r:   link_tower_typer;   r   ro   r(   tensorscaled_factorbetar=   r>   NotImplementedErrorrD   rG   r+   r,   r:   [  s    



zBridgeTowerLinkTower.__init__c                 C   st   | j dkr| || S | j dkr6| || j | S | j dkr^| |d| j  || j  S td| j  dd S )Nr   r   r   r   r   r   )r   r=   r   r   r   )rE   r"   Zcross_modal_hidden_statesrK   r+   r+   r,   rV   h  s    


zBridgeTowerLinkTower.forwardr$   r%   r&   r:   rV   rX   r+   r+   rG   r,   r   Z  s   r   c                       s4   e Zd Z fddZejejejdddZ  ZS )BridgeTowerSelfOutputc                    sB   t    t|j|j| _tj|j|jd| _t|j	| _
d S Nr5   )r9   r:   r   r@   r;   denser=   r>   Dropouthidden_dropout_probdropoutrD   rG   r+   r,   r:   u  s    
zBridgeTowerSelfOutput.__init__r"   input_tensorry   c                 C   s&   |  |}| |}| || }|S rR   r   r   r=   rE   r"   r   r+   r+   r,   rV   {  s    

zBridgeTowerSelfOutput.forwardr$   r%   r&   r:   r(   rW   rV   rX   r+   r+   rG   r,   r   t  s   r   c                       s0   e Zd Z fddZejejdddZ  ZS )BridgeTowerIntermediatec                    sB   t    t|j|j| _t|jt	r6t
|j | _n|j| _d S rR   )r9   r:   r   r@   r;   intermediate_sizer   
isinstance
hidden_actstrr   intermediate_act_fnrD   rG   r+   r,   r:     s
    
z BridgeTowerIntermediate.__init__r"   ry   c                 C   s   |  |}| |}|S rR   )r   r   rE   r"   r+   r+   r,   rV     s    

zBridgeTowerIntermediate.forwardr   r+   r+   rG   r,   r     s   r   c                       s4   e Zd Z fddZejejejdddZ  ZS )BridgeTowerOutputc                    sB   t    t|j|j| _tj|j|jd| _t	|j
| _d S r   )r9   r:   r   r@   r   r;   r   r=   r>   r   r   r   rD   rG   r+   r,   r:     s    
zBridgeTowerOutput.__init__r   c                 C   s&   |  |}| |}| || }|S rR   r   r   r+   r+   r,   rV     s    

zBridgeTowerOutput.forwardr   r+   r+   rG   r,   r     s   r   c                       s0   e Zd Z fddZejejdddZ  ZS )BridgeTowerPoolerc                    s*   t    t|j|j| _t | _d S rR   )r9   r:   r   r@   r;   r   ZTanh
activationrD   rG   r+   r,   r:     s    
zBridgeTowerPooler.__init__r   c                 C   s(   |d d df }|  |}| |}|S )Nr   )r   r   )rE   r"   Zfirst_token_tensorpooled_outputr+   r+   r,   rV     s    

zBridgeTowerPooler.forwardr   r+   r+   rG   r,   r     s   r   c                
       s   e Zd Zd fdd	ZejejdddZdejeej eej eej eej ee	e	ej   ee
 e	ej dd	d
Z  ZS )BridgeTowerSelfAttentionNc                    s   t    |j|j dkr>t|ds>td|j d|j d|j| _t|j|j | _| j| j | _t	
|j| j| _t	
|j| j| _t	
|j| j| _t	|j| _|pt|dd| _| jdks| jd	kr|j| _t	d
|j d | j| _|j| _d S )Nr   Zembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()position_embedding_typeabsoluterelative_keyrelative_key_queryri   r   )r9   r:   r;   num_attention_headshasattr
ValueErrorintattention_head_sizeall_head_sizer   r@   querykeyvaluer   Zattention_probs_dropout_probr   getattrr   max_position_embeddingsrs   distance_embedding
is_decoderrE   rF   r   rG   r+   r,   r:     s*    
  z!BridgeTowerSelfAttention.__init__)xry   c                 C   s6   |  d d | j| jf }||}|ddddS )Nrl   r   ri   r   r
   )sizer   r   viewr   )rE   r   Znew_x_shaper+   r+   r,   transpose_for_scores  s    
z-BridgeTowerSelfAttention.transpose_for_scoresFr"   rK   	head_maskencoder_hidden_statesencoder_attention_maskpast_key_valueoutput_attentionsry   c                 C   s  |  |}|d k	}	|	r4|d k	r4|d }
|d }|}n|	r^| | |}
| | |}|}nv|d k	r| | |}
| | |}tj|d |
gdd}
tj|d |gdd}n | | |}
| | |}| |}|d k	}| jr|
|f}t||
dd}| j	dks | j	dkr|j
d |
j
d  }}|r^tj|d tj|jd	dd}ntj|tj|jd	dd}tj|tj|jd	dd}|| }| || j d }|j|jd
}| j	dkrtd||}|| }n4| j	dkrtd||}td|
|}|| | }|t| j }|d k	r:|| }tjj|dd}| |}|d k	rf|| }t||}|dddd }| d d | jf }||}|r||fn|f}| jr||f }|S )Nr   r   ri   r{   rl   r   r   rL   rz   zbhld,lrd->bhlrzbhrd,lrd->bhlrr
   ) r   r   r   r   r(   r   r   matmulr   r   r}   r   longrN   r   rv   r   r   rO   rM   Zeinsummathsqrtr   r   
functionalZsoftmaxr   r   
contiguousr   r   )rE   r"   rK   r   r   r   r   r   Zmixed_query_layerZis_cross_attentionZ	key_layerZvalue_layerZquery_layer	use_cacheZattention_scoresZquery_lengthZ
key_lengthZposition_ids_lZposition_ids_rZdistanceZpositional_embeddingZrelative_position_scoresZrelative_position_scores_queryZrelative_position_scores_keyZattention_probsZcontext_layerZnew_context_layer_shapeoutputsr+   r+   r,   rV     sp    


 





z BridgeTowerSelfAttention.forward)N)NNNNNF)r$   r%   r&   r:   r(   rW   r   r   r)   r   rP   rV   rX   r+   r+   rG   r,   r     s$         r   c                
       sv   e Zd Zd
 fdd	Zdd Zdejeej eej eej eej ee	e	ej   ee
 e	ej ddd	Z  ZS )BridgeTowerAttentionNc                    s.   t    t||d| _t|| _t | _d S )Nr   )r9   r:   r   rE   r   outputsetpruned_headsr   rG   r+   r,   r:   :  s    

zBridgeTowerAttention.__init__c                 C   s   t |dkrd S t|| jj| jj| j\}}t| jj|| j_t| jj|| j_t| jj	|| j_	t| j
j|dd| j
_| jjt | | j_| jj| jj | j_| j|| _d S )Nr   r   r{   )lenr   rE   r   r   r   r   r   r   r   r   r   r   union)rE   headsindexr+   r+   r,   prune_heads@  s       z BridgeTowerAttention.prune_headsFr   c              	   C   s<   |  |||||||}| |d |}	|	f|dd   }
|
S )Nr   r   )rE   r   )rE   r"   rK   r   r   r   r   r   Zself_outputsattention_outputr   r+   r+   r,   rV   R  s    
	zBridgeTowerAttention.forward)N)NNNNNF)r$   r%   r&   r:   r   r(   rW   r   r)   r   rP   rV   rX   r+   r+   rG   r,   r   9  s$         r   c                       s.   e Zd Z fddZd	ddZdd Z  ZS )
BridgeTowerBertCrossLayerc                    sT   t    |j| _d| _t|| _|j| _|j| _t|| _t	|| _
t|| _d S )Nr   )r9   r:   chunk_size_feed_forwardseq_len_dimr   rQ   r   add_cross_attentioncrossattentionr   intermediater   r   rD   rG   r+   r,   r:   k  s    



z"BridgeTowerBertCrossLayer.__init__NFc              	   C   sz   | j ||d |d d}|d }	|dd  }
| j|	||||||d}|d }	|
|dd  }
t| j| j| j|	}|f|
 }
|
S )N)rK   r   r   r   r   r   )rK   r   r   r   r   r   rl   )rQ   r   r   feed_forward_chunkr   r   )rE   r"   r   rK   r   r   r   r   self_attention_outputsr   r   cross_attention_outputslayer_outputr+   r+   r,   rV   v  s8    	   
z!BridgeTowerBertCrossLayer.forwardc                 C   s   |  |}| ||}|S rR   r   r   rE   r   Zintermediate_outputr   r+   r+   r,   r     s    
z,BridgeTowerBertCrossLayer.feed_forward_chunk)NNNNF)r$   r%   r&   r:   rV   r   rX   r+   r+   rG   r,   r   j  s        
,r   c                
       st   e Zd Z fddZd
ejeej eej eej eej eeeej   ee	 eej dddZ
dd	 Z  ZS )BridgeTowerTextLayerc                    sr   t    |j| _d| _t|| _|j| _|j| _| jrZ| jsLt|  dt|dd| _	t
|| _t|| _d S )Nr   z> should be used as a decoder model if cross attention is addedr   r   )r9   r:   r   r   r   rQ   r   r   r   r   r   r   r   r   rD   rG   r+   r,   r:     s    


zBridgeTowerTextLayer.__init__NFr   c              	   C   s  |d k	r|d d nd }| j |||||d}	|	d }
| jrP|	dd }|	d }n|	dd  }d }| jr|d k	rt| dstd|  d|d k	r|d	d  nd }| |
||||||}|d }
||dd  }|d }|| }t| j| j| j|
}|f| }| jr||f }|S )
Nri   r   r   r   r   rl   r   z'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`r   )	rQ   r   r   r   r   r   r   r   r   )rE   r"   rK   r   r   r   r   r   Zself_attn_past_key_valuer   r   r   Zpresent_key_valueZcross_attn_present_key_valueZcross_attn_past_key_valuer   r   r+   r+   r,   rV     sV    


	   

zBridgeTowerTextLayer.forwardc                 C   s   |  |}| ||}|S rR   r   r   r+   r+   r,   r     s    
z'BridgeTowerTextLayer.feed_forward_chunk)NNNNNF)r$   r%   r&   r:   r(   rW   r   r)   r   rP   rV   r   rX   r+   r+   rG   r,   r     s$         Ar   c                       s   e Zd Z fddZd	ejeej eej eej eej eeeej   ee	 ee	 ee	 ee	 e
eej ef dddZ  ZS )
BridgeTowerTextEncoderc                    s:   t     | _t fddt jD | _d| _d S )Nc                    s   g | ]}t  qS r+   )r   r[   r]   r+   r,   r^     s     z3BridgeTowerTextEncoder.__init__.<locals>.<listcomp>F)	r9   r:   rF   r   r`   ra   r_   rU   gradient_checkpointingrD   rG   r]   r,   r:      s    
 zBridgeTowerTextEncoder.__init__NFT)r"   rK   r   r   r   past_key_valuesr   r   output_hidden_statesreturn_dictry   c              	      st  |	rdnd } rdnd } r(| j jr(dnd }| jrJ| jrJ|rJtd d}|rRdnd }t| jD ]\}}|	rv||f }|d k	r|| nd }|d k	r|| nd | jr| jrև fdd}tj	j

|||||||}n|||||| }|d }|r||d f7 } r`||d f }| j jr`||d	 f }q`|	r@||f }|
sbtd
d |||||fD S t|||||dS )Nr+   zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fc                    s    fdd}|S )Nc                     s    | f S rR   r+   )inputs)moduler   r   r+   r,   custom_forward)  s    zUBridgeTowerTextEncoder.forward.<locals>.create_custom_forward.<locals>.custom_forwardr+   )r   r   r   )r   r,   create_custom_forward(  s    z=BridgeTowerTextEncoder.forward.<locals>.create_custom_forwardr   rl   r   ri   c                 s   s   | ]}|d k	r|V  qd S rR   r+   r\   vr+   r+   r,   	<genexpr>M  s   z1BridgeTowerTextEncoder.forward.<locals>.<genexpr>)last_hidden_stater   r"   r#   cross_attentions)rF   r   r   ZtrainingloggerZwarning_once	enumeraterU   r(   utils
checkpointtupler   )rE   r"   rK   r   r   r   r   r   r   r   r   all_hidden_statesall_self_attentionsZall_cross_attentionsZnext_decoder_cacheiZlayer_moduleZlayer_head_maskr   Zlayer_outputsr+   r   r,   rV     sv    
	

zBridgeTowerTextEncoder.forward)	NNNNNNFFT)r$   r%   r&   r:   r(   rW   r   r)   r   rP   r   r   rV   rX   r+   r+   rG   r,   r     s.   	         r   c                       s2   e Zd ZdZ fddZd
ddZdd	 Z  ZS )BridgeTowerTextEmbeddingszV
    Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
    c                    s   t    tj|j|j|jd| _t|j|j| _	t|j
|j| _tj|j|jd| _t|j| _t|dd| _| jdt|jddd | jd	tj| j tjd
dd |j| _tj|j|j| jd| _	d S )N)padding_idxr5   r   r   rj   rk   Frm   token_type_idsrz   )r9   r:   r   rs   
vocab_sizer;   Zpad_token_idword_embeddingsr   position_embeddingsZtype_vocab_sizetoken_type_embeddingsr=   r>   r   r   r   r   r   ru   r(   rv   rw   zerosrj   r   r   r  rD   rG   r+   r,   r:   h  s.    
      z"BridgeTowerTextEmbeddings.__init__Nr   c                 C   s   |d kr*|d k	r t || j|}n
| |}|d k	r<| }n| d d }|d }|d krt| dr| jd d d |f }||d |}	|	}ntj|tj	| j
jd}|d kr| |}| |}
||
 }| jdkr| |}||7 }| |}| |}|S )Nrl   r   r  r   rL   r   )"create_position_ids_from_input_idsr  &create_position_ids_from_inputs_embedsr   r   r  rw   r(   r  r   rj   rN   r  r  r   r  r=   r   )rE   	input_idsr  rj   inputs_embedspast_key_values_lengthinput_shape
seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr  r   r  r+   r+   r,   rV     s0    








z!BridgeTowerTextEmbeddings.forwardc                 C   sN   |  dd }|d }tj| jd || j d tj|jd}|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
        Nrl   r   rL   r   )r   r(   rv   r  r   rN   Z	unsqueezerw   )rE   r  r  Zsequence_lengthrj   r+   r+   r,   r    s    	   z@BridgeTowerTextEmbeddings.create_position_ids_from_inputs_embeds)NNNNr   )r$   r%   r&   r'   r:   rV   r  rX   r+   r+   rG   r,   r  b  s            
(r  c                 C   s6   |  | }tj|dd|| | }| | 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
    r   r{   )ner   r(   ZcumsumZtype_asr   )r  r  r  maskZincremental_indicesr+   r+   r,   r    s    r  c                   @   s0   e Zd ZdZeZdZdZddgZdZ	dd Z
d	S )
BridgeTowerPreTrainedModelz
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    bridgetowerFr   r3   r   c                 C   s  t |tr|jjjd d|jjj d  }|jjjd }d|jjj d }|jjjD ]z}tjj	|j
j|| jj d tjj	|j
jj|| jj d tjj	|jjj|| jj d tjj	|jjj|| jj d qVtjj	|jjj|| jj d tjj	|jjjj|| jj d n\t |tjtjtjfrF|jjj	dd| jj d n(t |tjrn|jj  |jjd t |tjr|jd k	r|jj  d S )Ng      ri   )stdg        g?)Zmeanr   r   )r   BridgeTowerVisionModelvisualr   r;   r_   rb   r   initZnormal_r<   Zin_proj_weightrF   Zinitializer_factorZout_projr~   rA   r6   r8   r   rp   rt   r@   rq   rs   datar=   rh   Zzero_Zfill_)rE   r   Zproj_stdZattn_stdZfc_stdrf   r+   r+   r,   _init_weights  s.     
 
z(BridgeTowerPreTrainedModel._init_weightsN)r$   r%   r&   r'   r   config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_skip_keys_device_placementr%  r+   r+   r+   r,   r    s   r  c                       s6   e Zd ZeZ fddZedd ZdddZ  Z	S )	r!  c                    s   t  | t|| _d S rR   )r9   r:   r   r"  rD   rG   r+   r,   r:     s    zBridgeTowerVisionModel.__init__c                 C   s   | j jjjjS rR   )r"  r   rr   r~   rM   rE   r+   r+   r,   rM     s    zBridgeTowerVisionModel.dtypeNc                 C   s   |  || j|S rR   )r"  typerM   )rE   imageZ
image_maskr+   r+   r,   rV     s    zBridgeTowerVisionModel.forward)N)
r$   r%   r&   r   r&  r:   propertyrM   rV   rX   r+   r+   rG   r,   r!    s
   
r!  c                       s   e Zd ZdZeZd fdd	Zdd Zdd Zd	d
 Z	de
ej e
ej e
ej e
ej e
ej e
ej e
ej e
ej e
eej  e
e e
e e
e e
e eeej ef dddZ  ZS )BridgeTowerTextModela*  

    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*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
    Kaiser and Illia Polosukhin.

    To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
    to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
    `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.

    .. _*Attention is all you need*: https://arxiv.org/abs/1706.03762

    Tc                    sD   t  | || _t|| _t|| _|r2t|nd | _| 	  d S rR   )
r9   r:   rF   r  r   r   encoderr   pooler	post_init)rE   rF   Zadd_pooling_layerrG   r+   r,   r:     s    

zBridgeTowerTextModel.__init__c                 C   s   | j jS rR   r   r  r'  r+   r+   r,   get_input_embeddings   s    z)BridgeTowerTextModel.get_input_embeddingsc                 C   s   || j _d S rR   r/  rE   r   r+   r+   r,   set_input_embeddings#  s    z)BridgeTowerTextModel.set_input_embeddingsc                 C   s*   |  D ]\}}| jj| j| qdS )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)rS   r,  rU   rQ   r   )rE   Zheads_to_prunerU   r   r+   r+   r,   _prune_heads&  s    z!BridgeTowerTextModel._prune_headsN)r  rK   r  rj   r   r  r   r   r   r   r   r   r   ry   c                 C   s^  |dk	r|n| j j}|dk	r |n| j j}|dk	r4|n| j j}| j jrZ|
dk	rP|
n| j j}
nd}
|dk	rx|dk	rxtdn@|dk	r| || | }n"|dk	r| dd }ntd|\}}|dk	r|j	n|j	}|	dk	r|	d d j
d nd}|dkrtj||| f|d}|dkrft| jd	rT| jjddd|f }|||}|}ntj|tj|d
}| ||}| j jr|dk	r| \}}}||f}|dkrtj||d}| |}nd}| || j j}| j|||||d}| j||||||	|
|||d
}|d }| jdk	r$| |nd}|sB||f|dd  S t|||j|j|j|jdS )a  
        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *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` of shape `(batch_size, sequence_length)`, *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 (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
            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 timerl   z5You have to specify either input_ids or inputs_embedsr   ri   rN   r  rL   )r  rj   r  r  r  )	rK   r   r   r   r   r   r   r   r   r   )r  r!   r   r"   r#   r  )rF   r   r   use_return_dictr   r   r   Z%warn_if_padding_and_no_attention_maskr   rN   r}   r(   onesr   r   r  rw   r  r   get_extended_attention_maskZinvert_attention_maskZget_head_maskr_   r,  r-  r   r   r"   r#   r  )rE   r  rK   r  rj   r   r  r   r   r   r   r   r   r   r  r   r  rN   r  r  r  Zextended_attention_maskZencoder_batch_sizeZencoder_sequence_lengthrT   Zencoder_hidden_shapeZencoder_extended_attention_maskZembedding_outputZencoder_outputsZsequence_outputr   r+   r+   r,   rV   /  s    $




zBridgeTowerTextModel.forward)T)NNNNNNNNNNNNN)r$   r%   r&   r'   r   r&  r:   r0  r2  r3  r   r(   rW   r   r)   rP   r   r   r   rV   rX   r+   r+   rG   r,   r+    sF                r+  zqThe bare BridgeTower Model transformer outputting BridgeTowerModelOutput object without any specific head on top.c                       s   e Zd Z fddZdd Zdd Zeeee	e
ddeej eej eej eej eej eej eej eej ee ee ee ee eej eeej e	f d	d
dZdd Z  ZS )BridgeTowerModelc                    s  t     | _ j j jrJtj j| _	tj j| _
nDt fddt jD | _	t fddt jD | _
td j| _t| _t| _js jr| jjjD ](}| jjjjj|j_| jjjjj|j_qtfddt jD | _tfddt jD | _t | _t | _ tj! j j"d| _#tj! j j"d| _$ j%rt& | _'t& | _(nHt fddt jd	 D | _'t fd
dt jd	 D | _(| )  d S )Nc                    s   g | ]}t j jqS r+   r   r@   r;   r[   )rF   text_configr+   r,   r^     s     z-BridgeTowerModel.__init__.<locals>.<listcomp>c                    s   g | ]}t j jqS r+   r9  r[   )rF   vision_configr+   r,   r^     s     ri   c                    s   g | ]}t  qS r+   r   r[   r:  r+   r,   r^     s     c                    s   g | ]}t  qS r+   r<  r[   r=  r+   r,   r^     s     r5   c                    s   g | ]}t  qS r+   r   r[   r]   r+   r,   r^     s     r   c                    s   g | ]}t  qS r+   r>  r[   r]   r+   r,   r^     s     )*r9   r:   rF   r;  r:  Z$share_cross_modal_transformer_layersr   r@   r;   cross_modal_text_transformcross_modal_image_transformr`   ra   r_   rs   r  r!  vision_modelr+  
text_modelr   Z"init_layernorm_from_vision_encoderr"  Zcross_modal_ln_separater   r~   r$  rh   cross_modal_image_layerscross_modal_text_layersr   cross_modal_image_poolercross_modal_text_poolerr=   r>   cross_modal_text_layernormcross_modal_image_layernormZshare_link_tower_layersr   cross_modal_text_link_towercross_modal_image_link_towerr.  )rE   rF   r   rG   )rF   r:  r;  r,   r:     sP    




zBridgeTowerModel.__init__c                 C   s
   | j  S rR   )rB  r0  r'  r+   r+   r,   r0    s    z%BridgeTowerModel.get_input_embeddingsc                 C   s   | j | d S rR   )rB  r2  r1  r+   r+   r,   r2    s    z%BridgeTowerModel.set_input_embeddingsoutput_typer&  N)r  rK   r  rx   
pixel_maskr   r  r1   image_token_type_idxr   r   r   labelsry   c           ,      C   sr  |
dk	r|
n| j j}
|dk	r |n| j j}|r0dnd}|r<dnd}|rHdnd}|rTdnd}|
r`dnd}|dk	rp|n| j j}|	r|	nd}	| }| jj|d}|r||f7 }|dkrtj|tj	|j
d}| j|||j
}t| jjj| j j d }| jjjd| D ]$}|||d }|r||f7 }q|dkrP| jj|| jj}n|ddd}|rn||f7 }| jjjjd| D ]}||}|r||f7 }q| jj|| jj}| |}| tjdtj	|j
d|}| || }| |}| tj d|	tj	|j
d|}|| }| !|}tj|d|dftj	|j
d}| j|| |j
}| j"d |||||
d	}|d } | j#d |||||
d	}!|!d }"|r|| |"ff7 }|
r||d |!d ff7 }d}#t$|t| jjjD ]}$| jjj|$ ||d }| jjjj|$ || jj}| | jj|| }| j%|# }%| j&|# }&|%| || | |}'|&||"|}(| j"|#d  |'|(|||
d	}|d } | j#|#d  |(|'|||
d	}!|!d }"|#d7 }#|r||f7 }||f7 }|| |"ff7 }|
r||d |!d ff7 }q| |" })}*| '|)|*}+|r>|||f}|s`t(d
d |)|*|+||fD S t)|)|*|+||dS )a1  
        output_hidden_states (`bool`, *optional*):
            If set to `True`, hidden states are returned as a list containing the hidden states of text, image, and
            cross-modal components respectively. i.e. `(hidden_states_text, hidden_states_image,
            hidden_states_cross_modal)` where each element is a list of the hidden states of the corresponding
            modality. `hidden_states_txt/img` are a list of tensors corresponding to unimodal hidden states and
            `hidden_states_cross_modal` is a list of tuples containing `cross_modal_text_hidden_states` and
            `cross_modal_image_hidden_states` of each brdige layer.
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels are currently not supported.
        Returns:

        Examples:

        ```python
        >>> from transformers import BridgeTowerProcessor, BridgeTowerModel
        >>> from PIL import Image
        >>> import requests

        >>> # prepare image and text
        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)
        >>> text = "hello world"
        >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base")
        >>> model = BridgeTowerModel.from_pretrained("BridgeTower/bridgetower-base")

        >>> inputs = processor(image, text, return_tensors="pt")
        >>> outputs = model(**inputs)
        >>> outputs.keys()
        odict_keys(['text_features', 'image_features', 'pooler_output'])
        ```Nr+   r   )r  rL   r   ri   r   )rK   r   r   c                 s   s   | ]}|d k	r|V  qd S rR   r+   r   r+   r+   r,   r     s   z+BridgeTowerModel.forward.<locals>.<genexpr>)r   r    r!   r"   r#   )*rF   r   r   r5  r   rB  r   r(   r6  r   rN   r7  rO   r   r,  rU   r_   rA  r"  r   r(  rM   r   r   rb   r   r?  r  r  	expand_asrG  r@  fullrH  rD  rC  ra   rI  rJ  get_cls_featuresr  r   ),rE   r  rK   r  rx   rM  r   r  r1   rN  r   r   r   rO  Zall_hidden_states_textZall_hidden_states_imageZall_hidden_states_crossr  r	  r  r0   Zextend_text_masksZsplit_indexrU   rf   image_embeds_with_lnZcross_modal_textZtext_token_type_embeddingsimage_token_type_embeddingsZcross_modal_imageZextend_image_masksZlayer_outputs_textZcross_text_featuresZlayer_outputs_imageZcross_image_featuresZlink_layer_indexr
  Ztext_link_towerZimage_link_towerZcross_text_features_Zcross_image_features_r   r    Zcls_featuresr+   r+   r,   rV     s
   1











zBridgeTowerModel.forwardc                 C   s&   |  |}| |}tj||gddS )Nrl   r{   )rF  rE  r(   r   )rE   r   r    Zcls_features_textZcls_features_imager+   r+   r,   rS    s    

z!BridgeTowerModel.get_cls_features)NNNNNNNNNNNNN)r$   r%   r&   r:   r0  r2  r   BRIDGETOWER_INPUTS_DOCSTRINGr   r   _CONFIG_FOR_DOCr   r(   
LongTensorr)   r   rP   r   r   rW   rV   rS  rX   r+   r+   rG   r,   r8    sH   8
              ar8  c                       s$   e Zd Z fddZdd Z  ZS )"BridgeTowerPredictionHeadTransformc                    sV   t    t|j|j| _t|jtr6t	|j | _
n|j| _
tj|j|jd| _d S r   )r9   r:   r   r@   r;   r   r   r   r   r   transform_act_fnr=   r>   rD   rG   r+   r,   r:     s    
z+BridgeTowerPredictionHeadTransform.__init__c                 C   s"   |  |}| |}| |}|S rR   )r   rZ  r=   r   r+   r+   r,   rV     s    


z*BridgeTowerPredictionHeadTransform.forwardr   r+   r+   rG   r,   rY    s   	rY  c                       s&   e Zd Zd fdd	Zdd Z  ZS )BridgeTowerMLMHeadNc                    s\   t    || _t|| _tj|j|jj	dd| _
tt|jj	| _|d k	rX|| j
_d S )NF)rh   )r9   r:   rF   rY  	transformr   r@   r;   r:  r  decoderro   r(   r  rh   r~   )rE   rF   r~   rG   r+   r,   r:     s    

zBridgeTowerMLMHead.__init__c                 C   s   |  |}| || j }|S rR   )r\  r]  rh   )rE   r   	mlm_scorer+   r+   r,   rV     s    
zBridgeTowerMLMHead.forward)Nr   r+   r+   rG   r,   r[    s   	r[  c                       s$   e Zd Z fddZdd Z  ZS )BridgeTowerITMHeadc                    s   t    t|d| _d S Nri   r9   r:   r   r@   fc)rE   r;   rG   r+   r,   r:     s    
zBridgeTowerITMHead.__init__c                 C   s   |  |}|S rR   rb  )rE   r   	itm_scorer+   r+   r,   rV     s    
zBridgeTowerITMHead.forwardr   r+   r+   rG   r,   r_     s   r_  z\
    BridgeTower Model with a language modeling head on top as done during pretraining.
    c                       s   e Zd ZdgZ fddZdd Zdd Zee	de
eed	deej eej eej eej eej eej eej eej ee ee ee eej eeeej f dddZ  ZS )BridgeTowerForMaskedLMzmlm_score.decoder.weightc                    s,   t  | t|| _t|| _|   d S rR   )r9   r:   r8  r  r[  r^  r.  rD   rG   r+   r,   r:     s    

zBridgeTowerForMaskedLM.__init__c                 C   s   | j jS rR   r^  r]  r'  r+   r+   r,   get_output_embeddings  s    z,BridgeTowerForMaskedLM.get_output_embeddingsc                 C   s   || j _d S rR   rf  )rE   Znew_embeddingsr+   r+   r,   set_output_embeddings  s    z,BridgeTowerForMaskedLM.set_output_embeddingszbatch_size, sequence_lengthrK  Nr  rK   r  rx   rM  r   r  r1   r   r   r   rO  ry   c                 C   s   |dk	r|n| j j}| j|||||||||	|
|d}| |rB|jn|d }d}|dk	rt }||j}||d| j j	j
|d}|st|}|dk	r|f| S |S t|||j|jdS )aQ  
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. 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 in `[0, ..., config.vocab_size]`
        Returns:

        Examples:

        ```python
        >>> from transformers import BridgeTowerProcessor, BridgeTowerForMaskedLM
        >>> from PIL import Image
        >>> import requests

        >>> url = "http://images.cocodataset.org/val2017/000000360943.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
        >>> text = "a <mask> looking out of the window"

        >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
        >>> model = BridgeTowerForMaskedLM.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")

        >>> # prepare inputs
        >>> encoding = processor(image, text, return_tensors="pt")

        >>> # forward pass
        >>> outputs = model(**encoding)

        >>> results = processor.decode(outputs.logits.argmax(dim=-1).squeeze(0).tolist())

        >>> print(results)
        .a cat looking out of the window.
        ```N
rK   r  rx   rM  r   r  r1   r   r   r   r   rl   r.   r/   r"   r#   )rF   r5  r  r^  r   r	   rO   rN   r   r:  r  r  r   r"   r#   )rE   r  rK   r  rx   rM  r   r  r1   r   r   r   rO  r   Z
mlm_logitsZmasked_lm_lossloss_fctr   r+   r+   r,   rV   "  s:    1zBridgeTowerForMaskedLM.forward)NNNNNNNNNNNN)r$   r%   r&   Z_tied_weights_keysr:   rg  rh  r   rV  formatr   r   rW  r   r(   rX  r)   rP   r   r   rV   rX   r+   r+   rG   r,   re  
  sB   	
            re  z
    BridgeTower Model transformer with a classifier head on top (a linear layer on top of the final hidden state of the
    [CLS] token) for image-to-text matching.
    c                       s   e Zd Z fddZeeeeedde	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e e	e e	e e	e
j eeee
j f dddZ  ZS )	#BridgeTowerForImageAndTextRetrievalc                    s2   t  | t|| _t|jd | _|   d S r`  )r9   r:   r8  r  r_  r;   rd  r.  rD   rG   r+   r,   r:   ~  s    
z,BridgeTowerForImageAndTextRetrieval.__init__rK  Nri  c                 C   s   |dk	r|n| j j}| j|||||||||	|
|d}|r>|jn|d }| |}d}|dk	rxt }||j}|||}|st|}|dk	r|f| S |S t	|||j
|jdS )a  
        labels (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*):
            Labels for computing the image-text matching loss. 0 means the pairs don't match and 1 means they match.
            The pairs with 0 will be skipped for calculation.
        Returns:

        Examples:

        ```python
        >>> from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
        >>> import requests
        >>> from PIL import Image

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)
        >>> texts = ["An image of two cats chilling on a couch", "A football player scoring a goal"]

        >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
        >>> model = BridgeTowerForImageAndTextRetrieval.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")

        >>> # forward pass
        >>> scores = dict()
        >>> for text in texts:
        ...     # prepare inputs
        ...     encoding = processor(image, text, return_tensors="pt")
        ...     outputs = model(**encoding)
        ...     scores[text] = outputs.logits[0, 1].item()
        ```Nrj  ri   rk  )rF   r5  r  r!   rd  r	   rO   rN   r  r   r"   r#   )rE   r  rK   r  rx   rM  r   r  r1   r   r   r   rO  r   r!   r/   Zitm_lossrl  r   r+   r+   r,   rV     s<    -

z+BridgeTowerForImageAndTextRetrieval.forward)NNNNNNNNNNNN)r$   r%   r&   r:   r   rV  r   r   rW  r   r(   rX  r)   rP   r   r   rV   rX   r+   r+   rG   r,   rn  v  s<   

            rn  c                       s$   e Zd Z fddZdd Z  ZS )BridgeTowerContrastiveHeadc                    s   t    t||| _d S rR   ra  )rE   r;   Z
embed_sizerG   r+   r,   r:     s    
z#BridgeTowerContrastiveHead.__init__c                 C   s   |  |}|S rR   rc  )rE   r   r+   r+   r,   rV     s    
z"BridgeTowerContrastiveHead.forwardr   r+   r+   rG   r,   ro    s   ro  zl
    BridgeTower Model with a image-text contrastive head on top computing image-text contrastive loss.
    c                       s   e Zd Z fddZeeeeedd	e	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e
j e	e e	e e	e e	e eeee
j f dddZ  ZS )
!BridgeTowerForContrastiveLearningc                    sl   t  | t|| _t|j|j| _t|j|j| _t|jd |j| _	t
t| jj| _|   d S r`  )r9   r:   r8  r  ro  r;   Zcontrastive_hidden_sizeitc_text_headitc_image_headitc_cross_modal_headr   ro   r(   r   rF   Zlogit_scale_init_valuelogit_scaler.  rD   rG   r+   r,   r:     s    
z*BridgeTowerForContrastiveLearning.__init__rK  NT)r  rK   r  rx   rM  r   r  r1   r   r   r   return_lossry   c           !      C   sD  |dk	r|n| j j}| j|||||||||	d|d}|r>|jn|d }|rP|jn|d \}}}|d }|d }| jjj|}| jt	j
ddt	j| jjjjd	|}| j|| }tjj| |ddd
ddf ddd}tjj| |ddd
ddf dddj|jd}tjj| |dddj|jd}t	j|||gdd}| j j|jd}t	|| | }t	|| | }t	|| | }d}|rt	jt||jd}tj||}tj||}tj||}|| | d }|s*||||f|dd  } |dk	r&|f|  S | S t||||||j|j dS )a  
        return_loss (`bool`, *optional*):
            Whether or not to return the contrastive loss.
        Returns:

        Examples:

        ```python
        >>> from transformers import BridgeTowerProcessor, BridgeTowerForContrastiveLearning
        >>> import requests
        >>> from PIL import Image
        >>> import torch

        >>> image_urls = [
        ...     "https://farm4.staticflickr.com/3395/3428278415_81c3e27f15_z.jpg",
        ...     "http://images.cocodataset.org/val2017/000000039769.jpg",
        ... ]
        >>> texts = ["two dogs in a car", "two cats sleeping on a couch"]
        >>> images = [Image.open(requests.get(url, stream=True).raw) for url in image_urls]

        >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")
        >>> model = BridgeTowerForContrastiveLearning.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")

        >>> inputs = processor(images, texts, padding=True, return_tensors="pt")
        >>> loss = model(**inputs, return_loss=True).loss

        >>> inputs = processor(images, texts[::-1], padding=True, return_tensors="pt")
        >>> loss_swapped = model(**inputs, return_loss=True).loss

        >>> print("Loss", round(loss.item(), 4))
        Loss 0.0019

        >>> print("Loss with swapped images", round(loss_swapped.item(), 4))
        Loss with swapped images 2.126
        ```NTrj  ri   r
   rl   rP  r   rL   r   )r|   pr4  r   r{   g      @)r.   r/   r0   r1   r2   r"   r#   )!rF   r5  r  r!   r"   rA  r"  r   r  r(   rR  r   r~   rN   rQ  r@  r   r   	normalizerq  rr  rO   rs  r   rt  expr   trv   r   Zcross_entropyr-   r#   )!rE   r  rK   r  rx   rM  r   r  r1   r   r   r   ru  r   r!   Zhidden_states_txtZhidden_states_imgZhidden_states_cross_modalr0   rT  rU  r2   r/   rt  Zlogits_text_to_imageZlogits_text_to_crossZlogits_image_to_crossZitc_lossrO  Ztext_to_image_lossZtext_to_cross_lossZimage_to_cross_lossr   r+   r+   r,   rV     sr    4**z)BridgeTowerForContrastiveLearning.forward)NNNNNNNNNTNN)r$   r%   r&   r:   r   rV  r   r-   rW  r   r(   rX  r)   rP   r   r   rV   rX   r+   r+   rG   r,   rp    s<   
            rp  )r   )Or'   r   collectionsr   dataclassesr   typingr   r   r   r   r(   Ztorch.utils.checkpointr   Ztorch.nnr	   Zactivationsr   r   Zmodeling_outputsr   r   r   r   r   Zmodeling_utilsr   r   Zpytorch_utilsr   r   r  r   r   r   r   Zconfiguration_bridgetowerr   r   r   Z
get_loggerr$   r  rW  Z_CHECKPOINT_FOR_DOCZ_TOKENIZER_FOR_DOCZ)BRIDGETOWER_PRETRAINED_MODEL_ARCHIVE_LISTZBRIDGETOWER_START_DOCSTRINGrV  r   r-   Moduler3   rY   rg   r   r   r   r   r   r   r   r   r   r   r   r  r  r  r!  r+  r8  rY  r[  r_  re  rn  ro  rp  r+   r+   r+   r,   <module>   s   
<!,#1 1>WcZ
' 1  *
f_
