U
    9%e                    @   s  d Z ddl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 dd	lmZ dd
lmZmZmZmZ ddlmZ ddlmZmZmZ ddlmZmZm Z m!Z! ddl"m#Z#m$Z$ e %e&Z'dZ(dZ)dgZ*eG dd deZ+dCe	j,e	j-ee. dddZ/G dd dej0Z1G dd dej0Z2G dd dej0Z3G dd dej0Z4G dd  d ej0Z5G d!d" d"ej0Z6G d#d$ d$ej0Z7G d%d& d&ej0Z8G d'd( d(eZ9d)Z:d*Z;G d+d, d,ej0Z<G d-d. d.ej0Z=G d/d0 d0ej0Z>G d1d2 d2ej0Z?G d3d4 d4ej0Z@d5ZAG d6d7 d7ej0ZBed8e:G d9d: d:e9ZCG d;d< d<ej0ZDed=e:G d>d? d?e9ZEed@e:G dAdB dBe9ZFdS )DzPyTorch GIT model.    N)	dataclass)ListOptionalTupleUnion)nn)CrossEntropyLoss   )ACT2FN)ModelOutput)BaseModelOutputBaseModelOutputWithPastBaseModelOutputWithPoolingCausalLMOutputWithPast)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   )	GitConfigGitVisionConfigzmicrosoft/git-baser   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S )GitVisionModelOutputa  
    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.

    Args:
        image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
            The image embeddings obtained by applying the projection layer to the pooler_output.
        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        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image_embedslast_hidden_statehidden_states
attentions)__name__
__module____qualname____doc__r   r   torchFloatTensor__annotations__r   r   r   r    r'   r'   c/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/git/modeling_git.pyr   5   s
   
r   )maskdtypetgt_lenc                 C   sj   |   \}}|dk	r|n|}| ddddddf |d|||}d| }||tjt|jS )z_
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    Nr         ?)sizeexpandtomasked_fillr$   boolZfinfomin)r)   r*   r+   bszsrc_lenZexpanded_maskZinverted_maskr'   r'   r(   _expand_maskT   s
    *r5   c                       sL   e Zd ZdZ fddZd	eej eej eej e	ej
dddZ  ZS )
GitEmbeddingsz;Construct the embeddings from word and position embeddings.c                    s   t    tj|j|j|jd| _t|j|j| _	tj
|j|jd| _
t|j| _t|dd| _| jdt|jddd d S )	N)padding_idxZepsposition_embedding_typeabsoluteposition_idsr   F
persistent)super__init__r   	Embedding
vocab_sizehidden_sizeZpad_token_idword_embeddingsmax_position_embeddingsposition_embeddings	LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutgetattrr9   register_bufferr$   aranger.   selfconfig	__class__r'   r(   rA   e   s    
  zGitEmbeddings.__init__Nr   )	input_idsr;   inputs_embedspast_key_values_lengthreturnc           	      C   s   |d k	r|  }n|  d d }|d }|d krL| jd d ||| f }|d kr`| |}n|}| jdkr| |}||7 }| |}| |}|S )Nr=   r   r:   )r-   r;   rE   r9   rG   rH   rL   )	rQ   rU   r;   rV   rW   input_shape
seq_length
embeddingsrG   r'   r'   r(   forwardt   s    




zGitEmbeddings.forward)NNNr   )r    r!   r"   r#   rA   r   r$   Z
LongTensorr%   intTensorr\   __classcell__r'   r'   rS   r(   r6   b   s       r6   c                	       sx   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	e	ej   ee
 ee
 e	ej dd	d
Z  ZS )GitSelfAttentionNc                    s2  t    |j|j dkr>t|ds>td|j d|j d|j| _t|j|j | _| j| j | _t|j	j
|j	j d d | _|jd k	r|  j|j9  _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| _d S )Nr   Zembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()   r   r9   r:   relative_keyrelative_key_query)r@   rA   rD   num_attention_headshasattr
ValueErrorr]   attention_head_sizeall_head_sizevision_config
image_size
patch_sizeimage_patch_tokensnum_image_with_embeddingr   LinearquerykeyvaluerJ   Zattention_probs_dropout_probrL   rM   r9   rF   rB   distance_embeddingrQ   rR   r9   rS   r'   r(   rA      s.    

  zGitSelfAttention.__init__)xrX   c                 C   s6   |  d d | j| jf }||}|ddddS )Nr=   r   rb   r   r	   )r-   re   rh   viewpermute)rQ   ru   Znew_x_shaper'   r'   r(   transpose_for_scores   s    
z%GitSelfAttention.transpose_for_scoresFr   attention_mask	head_maskpast_key_valueoutput_attentionspixel_values_presentrX   c              	   C   s*  |  |}|r| jnd}|d k	r| | |}	| | |}
tj|	d d d d d |d d f |d |	d d d d dd d d f gdd}	tj|
d d d d d |d d f |d |
d d d d dd d d f gdd}
n | | |}	| | |}
| |}|d k	}|	d d d d |d d d f |
d d d d |d d d f f}t||	dd}| j	dks| j	dkrx|j
d |	j
d  }}|rt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rDtd||}|| }n4| j	dkrxtd||}td|	|}|| | }|t| j }|d k	r|| }tjj|dd}| |}|d k	r|| }t||
}|dddd }| d d | jf }||}|r||fn|f}||f }|S )Nr   r=   rb   dimr   rc   rd   r*   devicer*   zbhld,lrd->bhlrzbhrd,lrd->bhlrr	   ) rp   rm   rx   rq   rr   r$   catmatmul	transposer9   shapetensorlongr   rv   rO   rs   rF   r/   r*   Zeinsummathsqrtrh   r   
functionalsoftmaxrL   rw   
contiguousr-   ri   )rQ   r   rz   r{   r|   r}   r~   Zmixed_query_layercutoffZ	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(   r\      sf    	
PD 

 





zGitSelfAttention.forward)N)NNNFF)r    r!   r"   rA   r$   r^   rx   r   r%   r   r1   r\   r_   r'   r'   rS   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 )GitSelfOutputc                    sB   t    t|j|j| _tj|j|jd| _t|j	| _
d S Nr8   )r@   rA   r   ro   rD   denserH   rI   rJ   rK   rL   rP   rS   r'   r(   rA     s    
zGitSelfOutput.__init__r   input_tensorrX   c                 C   s&   |  |}| |}| || }|S Nr   rL   rH   rQ   r   r   r'   r'   r(   r\     s    

zGitSelfOutput.forwardr    r!   r"   rA   r$   r^   r\   r_   r'   r'   rS   r(   r     s   r   c                	       sl   e Zd Zd
 fdd	Zdd Zdejeej eej ee	e	ej   ee
 ee
 e	ej ddd	Z  ZS )GitAttentionNc                    s.   t    t||d| _t|| _t | _d S )N)r9   )r@   rA   r`   rQ   r   outputsetpruned_headsrt   rS   r'   r(   rA      s    

zGitAttention.__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   rQ   re   rh   r   r   rp   rq   rr   r   r   ri   union)rQ   headsindexr'   r'   r(   prune_heads'  s       zGitAttention.prune_headsFry   c           
      C   s:   |  ||||||}| |d |}|f|dd   }	|	S )Nr   r   )rQ   r   )
rQ   r   rz   r{   r|   r}   r~   Zself_outputsattention_outputr   r'   r'   r(   r\   9  s    	zGitAttention.forward)N)NNNFF)r    r!   r"   rA   r   r$   r^   r   r%   r   r1   r\   r_   r'   r'   rS   r(   r     s         r   c                       s0   e Zd Z fddZejejdddZ  ZS )GitIntermediatec                    sB   t    t|j|j| _t|jt	r6t
|j | _n|j| _d S r   )r@   rA   r   ro   rD   intermediate_sizer   
isinstance
hidden_actstrr
   intermediate_act_fnrP   rS   r'   r(   rA   Q  s
    
zGitIntermediate.__init__r   rX   c                 C   s   |  |}| |}|S r   )r   r   rQ   r   r'   r'   r(   r\   Y  s    

zGitIntermediate.forwardr   r'   r'   rS   r(   r   P  s   r   c                       s4   e Zd Z fddZejejejdddZ  ZS )	GitOutputc                    sB   t    t|j|j| _tj|j|jd| _t	|j
| _d S r   )r@   rA   r   ro   r   rD   r   rH   rI   rJ   rK   rL   rP   rS   r'   r(   rA   a  s    
zGitOutput.__init__r   c                 C   s&   |  |}| |}| || }|S r   r   r   r'   r'   r(   r\   g  s    

zGitOutput.forwardr   r'   r'   rS   r(   r   `  s   r   c                	       sj   e Zd Z fddZd
ejeej eej eeeej   ee	 ee	 eej dddZ
dd	 Z  ZS )GitLayerc                    s:   t    |j| _d| _t|| _t|| _t|| _	d S )Nr   )
r@   rA   chunk_size_feed_forwardseq_len_dimr   	attentionr   intermediater   r   rP   rS   r'   r(   rA   o  s    


zGitLayer.__init__NFry   c                 C   sv   |d k	r|d d nd }| j ||||||d}|d }	|dd }
|d }t| j| j| j|	}|f|
 }
|
|f }
|
S )Nrb   )r}   r|   r~   r   r   r=   )r   r   feed_forward_chunkr   r   )rQ   r   rz   r{   r|   r}   r~   Zself_attn_past_key_valueZself_attention_outputsr   r   Zpresent_key_valuelayer_outputr'   r'   r(   r\   w  s*    
   

zGitLayer.forwardc                 C   s   |  |}| ||}|S r   )r   r   )rQ   r   Zintermediate_outputr   r'   r'   r(   r     s    
zGitLayer.feed_forward_chunk)NNNFF)r    r!   r"   rA   r$   r^   r   r%   r   r1   r\   r   r_   r'   r'   rS   r(   r   n  s         #r   c                       s|   e Zd Z fddZd	ejeej eej eeeej   ee	 ee	 ee	 ee	 ee	 e
eej ef d
ddZ  ZS )

GitEncoderc                    s:   t     | _t fddt jD | _d| _d S )Nc                    s   g | ]}t  qS r'   )r   .0_rR   r'   r(   
<listcomp>  s     z'GitEncoder.__init__.<locals>.<listcomp>F)	r@   rA   rR   r   
ModuleListrangenum_hidden_layerslayergradient_checkpointingrP   rS   r   r(   rA     s    
 zGitEncoder.__init__NFT)
r   rz   r{   past_key_valuesr   r}   output_hidden_statesr~   return_dictrX   c
                    s>  | j r| jr|rtd d}|r&dnd }
 r2dnd }|r>dnd }t| jD ]\}}|rb|
|f }
|d k	rr|| nd }|d k	r|| nd | j r| jr fdd}tjj|||||}n|||| |}|d }|r||d f7 } rL||d f }qL|r|
|f }
|	s.t	d	d
 |||
|fD S t
|||
|dS )NzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fr'   c                    s    fdd}|S )Nc                     s    | f S r   r'   inputs)moduler}   r|   r'   r(   custom_forward  s    zIGitEncoder.forward.<locals>.create_custom_forward.<locals>.custom_forwardr'   r   r   r}   r|   r   r(   create_custom_forward  s    z1GitEncoder.forward.<locals>.create_custom_forwardr   r=   r   c                 s   s   | ]}|d k	r|V  qd S r   r'   r   vr'   r'   r(   	<genexpr>  s   z%GitEncoder.forward.<locals>.<genexpr>r   r   r   r   )r   trainingloggerZwarning_once	enumerater   r$   utils
checkpointtupler   )rQ   r   rz   r{   r   r   r}   r   r~   r   Zall_hidden_statesZall_self_attentionsZnext_decoder_cacheiZlayer_moduleZlayer_head_maskr   layer_outputsr'   r   r(   r\     sf    
	

zGitEncoder.forward)NNNNFFFT)r    r!   r"   rA   r$   r^   r   r%   r   r1   r   r   r\   r_   r'   r'   rS   r(   r     s*   	        r   c                   @   s.   e Zd ZdZeZdZdZdd Zd
ddZ	d	S )GitPreTrainedModelz
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    gitTc                 C   s   t |trRtjj|jd| jjd tjj|jj	| jjd tjj|j
j	| jjd t |tjr|j	jjd| jjd |jdk	r|jj  nft |tjr|j	jjd| jjd |jdk	r|j	j|j   n&t |tjr|jj  |j	jd dS )zInitialize the weightsg        )Zmeanstd)r   Nr,   )r   GitVisionEmbeddingsr   initZnormal_class_embeddingrR   Zinitializer_rangepatch_embeddingweightposition_embeddingro   databiasZzero_rB   r7   rH   Zfill_)rQ   r   r'   r'   r(   _init_weights  s    


z GitPreTrainedModel._init_weightsFc                 C   s   t |ttfr||_d S r   )r   r   GitVisionEncoderr   )rQ   r   rr   r'   r'   r(   _set_gradient_checkpointing  s    z.GitPreTrainedModel._set_gradient_checkpointingN)F)
r    r!   r"   r#   r   config_classZbase_model_prefixZsupports_gradient_checkpointingr   r   r'   r'   r'   r(   r     s   r   a=  

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also 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 ([`GitConfig`]): 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)

        position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
            config.max_position_embeddings - 1]`.

            [What are position IDs?](../glossary#position-ids)

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

        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.
        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                       s6   e Zd Zed fddZejejdddZ  Z	S )r   r   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strider   rb   r   r;   r<   r>   )r@   rA   rR   rD   	embed_dimrk   rl   r   	Parameterr$   Zrandnr   ZConv2dZnum_channelsr   Znum_patchesZnum_positionsrB   r   rN   rO   r.   rP   rS   r'   r(   rA   _  s"    
zGitVisionEmbeddings.__init__)pixel_valuesrX   c                 C   sn   |j d }| jjj}| |j|d}|ddd}| j|dd}t	j
||gdd}|| | j }|S )Nr   r   rb   r   r=   r   )r   r   r   r*   r/   flattenr   r   r.   r$   r   r   r;   )rQ   r   Z
batch_sizeZtarget_dtypeZpatch_embedsZclass_embedsr[   r'   r'   r(   r\   u  s    

zGitVisionEmbeddings.forward)
r    r!   r"   r   rA   r$   r%   r^   r\   r_   r'   r'   rS   r(   r   ^  s   r   c                       s0   e Zd Z fddZejejdddZ  ZS )GitVisionMLPc                    sD   t    || _t|j | _t|j|j	| _
t|j	|j| _d S r   )r@   rA   rR   r
   r   activation_fnr   ro   rD   r   fc1fc2rP   rS   r'   r(   rA     s
    
zGitVisionMLP.__init__r   c                 C   s"   |  |}| |}| |}|S r   )r   r   r   r   r'   r'   r(   r\     s    


zGitVisionMLP.forwardr   r'   r'   rS   r(   r     s   r   c                       sz   e Zd ZdZ fddZejeedddZdeje	ej e	ej e	e
 eeje	ej e	eej  f d	d
dZ  ZS )GitVisionAttentionz=Multi-headed attention from 'Attention Is All You Need' paperc                    s   t    || _|j| _|j| _| j| j | _| j| j | jkrZtd| j d| j d| jd | _	|j
| _t| j| j| _t| j| j| _t| j| j| _t| j| j| _d S )Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).g      )r@   rA   rR   rD   r   re   	num_headshead_dimrg   scaleZattention_dropoutrL   r   ro   k_projv_projq_projout_projrP   rS   r'   r(   rA     s    
zGitVisionAttention.__init__)r   seq_lenr3   c                 C   s    | ||| j| jdd S )Nr   rb   )rv   r   r   r   r   )rQ   r   r   r3   r'   r'   r(   _shape  s    zGitVisionAttention._shapeNFr   rz   causal_attention_maskr}   rX   c                 C   s  |  \}}}| || j }| | |d|}	| | |d|}
|| j d| jf}| |||j| }|	j| }	|
j| }
|	 d}t	
||	dd}|  || j ||fkrtd|| j ||f d|   |dk	rD|  |d||fkrtd|d||f d|   ||| j||| }||| j ||}|dk	r|  |d||fkrtd|d||f d|   ||| j||| }||| j ||}tjj|dd}|r||| j||}||| j ||}nd}tjj|| j| jd	}t	
||
}|  || j || jfkrRtd
|| j|| jf d|   ||| j|| j}|dd}||||}| |}||fS )z#Input shape: Batch x Time x Channelr=   r   rb   z$Attention weights should be of size z	, but is Nz!Attention mask should be of size r   )pr   z `attn_output` should be of size )r-   r   r   r   r   r   r   r   rv   r$   Zbmmr   rg   r   r   r   rL   r   Zreshaper   )rQ   r   rz   r   r}   r3   r+   r   Zquery_statesZ
key_statesZvalue_statesZ
proj_shaper4   attn_weightsZattn_weights_reshapedZ
attn_probsZattn_outputr'   r'   r(   r\     sX    	





zGitVisionAttention.forward)NNF)r    r!   r"   r#   rA   r$   r^   r]   r   r   r1   r   r\   r_   r'   r'   rS   r(   r     s      r   c                       sJ   e Zd Zed fddZdejejejee e	ej
 dddZ  ZS )	GitVisionEncoderLayerr   c                    sR   t    |j| _t|| _tj| j|jd| _	t
|| _tj| j|jd| _d S r   )r@   rA   rD   r   r   	self_attnr   rH   rI   layer_norm1r   mlplayer_norm2rP   rS   r'   r(   rA     s    


zGitVisionEncoderLayer.__init__Fr   c                 C   sd   |}|  |}| j||||d\}}|| }|}| |}| |}|| }|f}|r`||f7 }|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.
        )r   rz   r   r}   )r   r   r  r   )rQ   r   rz   r   r}   Zresidualr   r   r'   r'   r(   r\     s"    




zGitVisionEncoderLayer.forward)F)r    r!   r"   r   rA   r$   r^   r   r1   r   r%   r\   r_   r'   r'   rS   r(   r     s    r   c                	       s`   e Zd ZdZed fddZd	eej eej ee	 ee	 ee	 e
eef dddZ  ZS )
r   z
    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
    [`GitVisionEncoderLayer`].

    Args:
        config: GitVisionConfig
    r   c                    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     z-GitVisionEncoder.__init__.<locals>.<listcomp>F)	r@   rA   rR   r   r   r   r   layersr   rP   rS   r   r(   rA   7  s    
 zGitVisionEncoder.__init__N)rz   r   r}   r   r   rX   c                    s   dk	r n| j j |dk	r |n| j j}|dk	r4|n| j j}|rDdnd} rPdnd}|}	t| jD ]r\}
}|rx||	f }| jr| jr fdd}tj	j

|||	||}n||	|| d}|d }	 rb||d f }qb|r||	f }|stdd	 |	||fD S t|	||d
S )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'   c                    s    fdd}|S )Nc                     s    | f S r   r'   r   )r   r}   r'   r(   r   s  s    zOGitVisionEncoder.forward.<locals>.create_custom_forward.<locals>.custom_forwardr'   r   r}   r   r(   r   r  s    z7GitVisionEncoder.forward.<locals>.create_custom_forwardr  r   r   c                 s   s   | ]}|d k	r|V  qd S r   r'   r   r'   r'   r(   r     s      z+GitVisionEncoder.forward.<locals>.<genexpr>r   r   r   )rR   r}   r   use_return_dictr   r  r   r   r$   r   r   r   r   )rQ   rV   rz   r   r}   r   r   Zencoder_statesZall_attentionsr   idxZencoder_layerr   r   r'   r  r(   r\   =  sH    &

  zGitVisionEncoder.forward)NNNNN)r    r!   r"   r#   r   rA   r   r$   r^   r1   r   r   r   r\   r_   r'   r'   rS   r(   r   .  s   	     
r   aE  
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
        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ed fddZeeeeedd	e	e
j e	e e	e e	e eeef dddZ  ZS )
GitVisionTransformerr   c                    sR   t    || _|j}t|| _tj||jd| _	t
|| _tj||jd| _d S r   )r@   rA   rR   rD   r   r[   r   rH   rI   pre_layrnormr   encoderpost_layernorm)rQ   rR   r   rS   r'   r(   rA     s    


zGitVisionTransformer.__init__output_typer   Nr   r}   r   r   rX   c                 C   s   |dk	r|n| j j}|dk	r |n| j j}|dk	r4|n| j j}|dkrLtd| |}| |}| j||||d}|d }| |}|s|f|dd  S t	||j
|jdS )z
        Returns:

        Nz You have to specify pixel_values)rV   r}   r   r   r   r   r  )rR   r}   r   r  rg   r[   r  r	  r
  r   r   r   )rQ   r   r}   r   r   r   encoder_outputsr   r'   r'   r(   r\     s.    


zGitVisionTransformer.forward)NNNN)r    r!   r"   r   rA   r   GIT_VISION_INPUTS_DOCSTRINGr   r   r   r$   r%   r1   r   r   r\   r_   r'   r'   rS   r(   r    s   

    
r  zOThe vision model from CLIP, used in GIT, without any head or projection on top.c                
       s   e Zd ZeZdZed fddZejdddZ	e
eeeeddeej ee ee ee eeef d
ddZ  ZS )GitVisionModelr   r   c                    s"   t  | t|| _|   d S r   )r@   rA   r  vision_model	post_initrP   rS   r'   r(   rA     s    
zGitVisionModel.__init__)rX   c                 C   s
   | j jjS r   )r  r[   r   rQ   r'   r'   r(   get_input_embeddings  s    z#GitVisionModel.get_input_embeddingsr  Nr  c                 C   s&   |dk	r|n| j j}| j||||dS )a  
        Returns:

        Examples:

        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, GitVisionModel

        >>> processor = AutoProcessor.from_pretrained("microsoft/git-base")
        >>> model = GitVisionModel.from_pretrained("microsoft/git-base")

        >>> 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
        ```N)r   r}   r   r   )rR   r  r  )rQ   r   r}   r   r   r'   r'   r(   r\     s    zGitVisionModel.forward)NNNN)r    r!   r"   r   r   Zmain_input_namerA   r   Moduler  r   r  r   r   r   r$   r%   r1   r   r   r\   r_   r'   r'   rS   r(   r    s"   
    
r  c                       s6   e Zd Zed fddZejejdddZ  ZS )GitProjectionr   c                    s@   t    || _tt|jj|jtj|j|jj	d| _
d S r   )r@   rA   rR   r   Z
Sequentialro   rj   rD   rH   rI   visual_projectionrP   rS   r'   r(   rA     s    
zGitProjection.__init__)r[   rX   c                 C   s
   |  |S r   )r  )rQ   r[   r'   r'   r(   r\   "  s    zGitProjection.forward)	r    r!   r"   r   rA   r$   r^   r\   r_   r'   r'   rS   r(   r    s   r  zThe bare GIT Model transformer consisting of a CLIP image encoder and text decoder outputting raw hidden-states without any specific head on top.c                       s   e Zd Z fddZdd Zdd Zdd Zeej	ej
ejd	d
dZd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ej  ee ee ee ee eeej ef dddZ  ZS )GitModelc                    sr   t     | _t | _t j| _t | _	t
 | _ jd k	rft fddt jD | _|   d S )Nc                 3   s&   | ]}t td d  jjV  qdS )r   N)r   r   r$   zerosrj   rD   r   r   r'   r(   r   7  s   z$GitModel.__init__.<locals>.<genexpr>)r@   rA   rR   r6   r[   r  rj   image_encoderr   r	  r  r  rn   r   ZParameterListr   img_temperal_embeddingr  rP   rS   r   r(   rA   ,  s    




zGitModel.__init__c                 C   s   | j jS r   r[   rE   r  r'   r'   r(   r  ?  s    zGitModel.get_input_embeddingsc                 C   s   || j _d S r   r  )rQ   rr   r'   r'   r(   set_input_embeddingsB  s    zGitModel.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)itemsr	  r   r   r   )rQ   Zheads_to_pruner   r   r'   r'   r(   _prune_headsE  s    zGitModel._prune_heads)r-   r*   r   rX   c                 C   s4   t jt j||||ddd}||dktd}|S )Nr   r*   r   )Zdiagonal-inf)r$   ZtriuZonesr0   float)rQ   r-   r*   r   r)   r'   r'   r(   _generate_future_maskM  s    zGitModel._generate_future_maskNc                 C   s  |j d }|j d }|j}|j}	tj||f||	d}
tj||| ftd|j|	d}tj||f|	|jd}|dkrtj|j d |j d | f|	|jd}tj|
|fdd}tj|||	fdd}tj||fddd d d f }|d kr
tj|j d |j d fd|d}|jtj	kr t
d	tj||jd
}td||< ||j d || || | f}| }|d d d d d |f }|d d d d d f }|| |d d d d d |f< |d d d d d d d f }|S )Nr   r   r!  r   r   r   F)Z
fill_valuer   z1Memory key padding mask must be a boolean tensor.r   )r   r   r*   r$   r  fullr"  r   r/   r1   rg   Z
zeros_liker.   clone)rQ   tgtmemorytgt_maskrW   Zmemory_key_padding_maskZnum_tgtZ
num_memoryr   r*   top_left	top_rightbottom_leftleftrightZfull_attention_maskZzero_negative_infinityZorigin_leftupdater'   r'   r(   create_attention_maskS  sP    



 zGitModel.create_attention_maskbatch_size, sequence_lengthr  )rU   rz   r;   r   r{   rV   r   r   r}   r   r   rX   c                 C   s   |	dk	r|	n| j j}	|
dk	r |
n| j j}
|dk	r4|n| j j}|dk	rH|n| j j}|dk	rj|dk	rjtdn@|dk	r| || | }n"|dk	r| dd }ntd|d }|dk	r|d d jd nd}| 	|| j j
}d}|dk	r|jdkr| |j}nz|jd	krzg }t|jd D ]B}| |dd|ddddf j}|| j| 7 }|| q&tj|dd
}ntd| |}| j||||d}|dkrtj|jd d|jd f|j|jd}||d|d dd}tj||fdd
}| ||j|j}| j||||d}|dk	rt||j|d d|j}|dkrv|dddd| dddf }n4|dddd|d  d|d  df  |7  < | j||||||	|
||dk	d	}|d }|s|f|dd  S t||j|j |j!dS )a  
        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`).

        Returns:

        Examples:

        ```python
        >>> from transformers import AutoProcessor, AutoModel
        >>> import requests
        >>> from PIL import Image

        >>> processor = AutoProcessor.from_pretrained("microsoft/git-base")
        >>> model = AutoModel.from_pretrained("microsoft/git-base")

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> text = "this is an image of two cats"

        >>> inputs = processor(text, images=image, return_tensors="pt")

        >>> outputs = model(**inputs)
        >>> last_hidden_state = outputs.last_hidden_state
        ```NzDYou cannot specify both input_ids and inputs_embeds at the same timer=   z5You have to specify either input_ids or inputs_embedsr   r   rb         r   z#pixel_values must be of rank 4 or 5)rU   r;   rV   rW   r   )r&  r'  r(  rW   )r+   )rz   r{   r   r   r}   r   r   r~   r   )"rR   r}   r   r   r  rg   Z%warn_if_padding_and_no_attention_maskr-   r   Zget_head_maskr   ndimr  r   r   r  appendr$   r   r  r[   r  r*   r   repeatr#  r/  r5   r/   r	  r   r   r   r   )rQ   rU   rz   r;   r   r{   rV   r   r   r}   r   r   rY   rZ   rW   Zprojected_visual_featuresZvisual_featuresZ	frame_idxZvisual_features_frameZembedding_outputr   r(  Zcombined_attention_maskZexpanded_attn_maskr  sequence_outputr'   r'   r(   r\     s    1


$

  

$4zGitModel.forward)N)NNNNNNNNNNN)r    r!   r"   rA   r  r  r  r]   r$   r*   r   r^   r#  r/  r   GIT_INPUTS_DOCSTRINGformatr   r   _CONFIG_FOR_DOCr   r   r%   r1   r   r   r\   r_   r'   r'   rS   r(   r  &  sB   
2
           r  zVGIT Model with a `language modeling` head on top for autoregressive language modeling.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ej  ee ee ee ee eeej ef dddZdddZdd Z  ZS )GitForCausalLMzoutput.weightc                    s4   t  | t|| _t|j|j| _| 	  d S r   )
r@   rA   r  r   r   ro   rD   rC   r   r  rP   rS   r'   r(   rA   4  s    
zGitForCausalLM.__init__c                 C   s   | j S r   r   r  r'   r'   r(   get_output_embeddings=  s    z$GitForCausalLM.get_output_embeddingsc                 C   s
   || _ d S r   r;  )rQ   Znew_embeddingsr'   r'   r(   set_output_embeddings@  s    z$GitForCausalLM.set_output_embeddingsr0  r  N)rU   rz   r;   r   r{   rV   labelsr   r   r}   r   r   rX   c                 C   s  |dk	r|n| j j}|dk	r d}	| j||||||||	|
||d}|d }| |}d}|dk	r| jjjd jjj}|dd|dddf 	 }|ddddf 	 }t
 }||d| j j|d}|s|f|dd  }|dk	r|f| S |S t|||j|j|jdS )a$  
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *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 (`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`).

        Returns:

        Examples:

        Image captioning example:

        ```python
        >>> from transformers import AutoProcessor, AutoModelForCausalLM
        >>> import requests
        >>> from PIL import Image

        >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
        >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> pixel_values = processor(images=image, return_tensors="pt").pixel_values

        >>> generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
        >>> generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        >>> print(generated_caption)
        two cats sleeping on a pink blanket next to remotes.
        ```

        Visual question answering (VQA) example:

        ```python
        >>> from transformers import AutoProcessor, AutoModelForCausalLM
        >>> from huggingface_hub import hf_hub_download
        >>> from PIL import Image

        >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-textvqa")
        >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-textvqa")

        >>> file_path = hf_hub_download(repo_id="nielsr/textvqa-sample", filename="bus.png", repo_type="dataset")
        >>> image = Image.open(file_path).convert("RGB")

        >>> pixel_values = processor(images=image, return_tensors="pt").pixel_values

        >>> question = "what does the front of the bus say at the top?"

        >>> input_ids = processor(text=question, add_special_tokens=False).input_ids
        >>> input_ids = [processor.tokenizer.cls_token_id] + input_ids
        >>> input_ids = torch.tensor(input_ids).unsqueeze(0)

        >>> generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
        >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True))
        ['what does the front of the bus say at the top? special']
        ```

        Video captioning example:

        ```python
        >>> import av
        >>> import numpy as np
        >>> from PIL import Image
        >>> from huggingface_hub import hf_hub_download
        >>> from transformers import AutoProcessor, AutoModelForCausalLM

        >>> processor = AutoProcessor.from_pretrained("microsoft/git-base-vatex")
        >>> model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-vatex")

        >>> # set seed for reproducability
        >>> np.random.seed(45)


        >>> def read_video_pyav(container, indices):
        ...     '''
        ...     Decode the video with PyAV decoder.
        ...     Args:
        ...         container (`av.container.input.InputContainer`): PyAV container.
        ...         indices (`List[int]`): List of frame indices to decode.
        ...     Returns:
        ...         result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3).
        ...     '''
        ...     frames = []
        ...     container.seek(0)
        ...     start_index = indices[0]
        ...     end_index = indices[-1]
        ...     for i, frame in enumerate(container.decode(video=0)):
        ...         if i > end_index:
        ...             break
        ...         if i >= start_index and i in indices:
        ...             frames.append(frame)
        ...     return np.stack([x.to_ndarray(format="rgb24") for x in frames])


        >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
        ...     '''
        ...     Sample a given number of frame indices from the video.
        ...     Args:
        ...         clip_len (`int`): Total number of frames to sample.
        ...         frame_sample_rate (`int`): Sample every n-th frame.
        ...         seg_len (`int`): Maximum allowed index of sample's last frame.
        ...     Returns:
        ...         indices (`List[int]`): List of sampled frame indices
        ...     '''
        ...     converted_len = int(clip_len * frame_sample_rate)
        ...     end_idx = np.random.randint(converted_len, seg_len)
        ...     start_idx = end_idx - converted_len
        ...     indices = np.linspace(start_idx, end_idx, num=clip_len)
        ...     indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
        ...     return indices


        >>> # load video
        >>> file_path = hf_hub_download(
        ...     repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset"
        ... )
        >>> container = av.open(file_path)

        >>> # sample frames
        >>> num_frames = model.config.num_image_with_embedding
        >>> indices = sample_frame_indices(
        ...     clip_len=num_frames, frame_sample_rate=4, seg_len=container.streams.video[0].frames
        ... )
        >>> frames = read_video_pyav(container, indices)

        >>> pixel_values = processor(images=list(frames), return_tensors="pt").pixel_values

        >>> generated_ids = model.generate(pixel_values=pixel_values, max_length=50)

        >>> print("Generated caption:", processor.batch_decode(generated_ids, skip_special_tokens=True))
        Generated caption: ['a woman is sitting at a table and she is talking about the food she is holding.']
        ```
        NF)
rz   r;   r   r{   rV   r   r   r}   r   r   r   r=   r   )losslogitsr   r   r   )rR   r  r   r   r	  r   r   rQ   rm   r   r   rv   rC   r   r   r   r   )rQ   rU   rz   r;   r   r{   rV   r>  r   r   r}   r   r   r   r6  r@  r?  Znum_image_tokensZshifted_logitsZloss_fctr   r'   r'   r(   r\   C  sH      
zGitForCausalLM.forwardc                 K   sL   |d k	r|d d dd f }|j }|d kr4||}|||dd ||dS )Nr=   r   )rU   rz   r   r   r   )r   Znew_onesget)rQ   rU   r   rz   r   kwargsrY   r'   r'   r(   prepare_inputs_for_generation  s    

z,GitForCausalLM.prepare_inputs_for_generationc                    s.   d}|D ] }|t  fdd|D f7 }q|S )Nr'   c                 3   s"   | ]}| d  |jV  qdS )r   N)Zindex_selectr/   r   )r   Z
past_statebeam_idxr'   r(   r   $  s     z0GitForCausalLM._reorder_cache.<locals>.<genexpr>)r   )rQ   r   rE  Zreordered_pastZ
layer_pastr'   rD  r(   _reorder_cache   s    zGitForCausalLM._reorder_cache)NNNNNNNNNNNN)NNN)r    r!   r"   Z_tied_weights_keysrA   r<  r=  r   r7  r8  r   r   r9  r   r$   r^   r   r1   r   r   r\   rC  rF  r_   r'   r'   rS   r(   r:  .  sN   	
             I     
r:  )N)Gr#   r   dataclassesr   typingr   r   r   r   r$   Ztorch.utils.checkpointr   Ztorch.nnr   Zactivationsr
   Z
file_utilsr   Zmodeling_outputsr   r   r   r   Zmodeling_utilsr   Zpytorch_utilsr   r   r   r   r   r   r   r   Zconfiguration_gitr   r   Z
get_loggerr    r   Z_CHECKPOINT_FOR_DOCr9  Z!GIT_PRETRAINED_MODEL_ARCHIVE_LISTr   r^   r*   r]   r5   r  r6   r`   r   r   r   r   r   r   r   ZGIT_START_DOCSTRINGr7  r   r   r   r   r   r  r  r  r  r  r:  r'   r'   r'   r(   <module>   st   
0~22Z#1$i3g:6   