o
    Zh                     @   s  d Z ddlZddlmZ ddlmZmZmZ ddlZddl	Zddlm
Z
 ddlmZmZmZ ddlmZ dd	lmZmZ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$ e"%e&Z'da(dd Z)dd Z*dd Z+dd Z,G dd dej-j.Z/G dd dej-j.Z0G dd de
j1Z2G dd de
j1Z3G dd  d e
j1Z4G d!d" d"e
j1Z5G d#d$ d$e
j1Z6G d%d& d&e
j1Z7G d'd( d(e
j1Z8G d)d* d*e
j1Z9G d+d, d,e
j1Z:G d-d. d.e
j1Z;G d/d0 d0e
j1Z<eG d1d2 d2eZ=eG d3d4 d4e=Z>eG d5d6 d6e=Z?G d7d8 d8e
j1Z@ed9d:G d;d< d<e=ZAeG d=d> d>e=ZBeG d?d@ d@e=ZCeG dAdB dBe=ZDg dCZEdS )DzPyTorch YOSO model.    N)Path)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss   )ACT2FN)"BaseModelOutputWithCrossAttentionsMaskedLMOutputMultipleChoiceModelOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringis_ninja_availableis_torch_cuda_availablelogging   )
YosoConfigc                  C   s:   ddl m}  dd }|g d}| d|dd dd lad S )	Nr   )loadc                    s,   t t jjjd d   fdd| D S )NZkernelsyosoc                    s   g | ]} | qS  r   ).0fileZ
src_folderr   U/var/www/auris/lib/python3.10/site-packages/transformers/models/yoso/modeling_yoso.py
<listcomp>:       z:load_cuda_kernels.<locals>.append_root.<locals>.<listcomp>)r   __file__resolveparent)filesr   r!   r"   append_root8   s   z&load_cuda_kernels.<locals>.append_root)zfast_lsh_cumulation_torch.cppzfast_lsh_cumulation.cuzfast_lsh_cumulation_cuda.cufast_lsh_cumulationT)verbose)Ztorch.utils.cpp_extensionr   r*   lsh_cumulation)r   r)   Z	src_filesr   r   r"   load_cuda_kernels4   s
   r-   c                 C   sJ   t | trg }| D ]}| s| }|| q	|S |  s#|  } | S N)
isinstancelistZis_contiguous
contiguousappendZinput_tensorsouttensorr   r   r"   to_contiguousC   s   
r6   c                 C   sF   t | trg }| D ]}|tjj|ddd q	|S tjj| dddS )N   )pdim)r/   r0   r2   r   Z
functional	normalizer3   r   r   r"   r;   Q   s   
r;   c                 C   s   t |  dkrtdt | dkrtdtj| d| d|| | jd}dtj|| jd }t| || d| d||}t|||d|d||}|dk	 }|dk	 }	tj
|| dd	}
tj
|	| dd	}
|
	 |
	 fS )
Nr
   zQuery has incorrect size.zKey has incorrect size.r   r7   devicer   r8   r:   )lensize
ValueErrortorchZrandnr=   arangematmulreshapeintsum)querykeynum_hashZhash_lenZrmatZ	raise_powZquery_projectionZkey_projectionZquery_binaryZ
key_binaryZ
query_hashr   r   r"   hashing[   s   $$$rK   c                   @   $   e Zd Zedd Zedd ZdS )YosoCumulationc           
   
   C   s   |d }dt t ||ddtj  | }||d d d d d f  |d d d d d f  }t ||}	| |||||| || _|	S )Nhash_code_lenr   r8   )rB   acosrD   	transposemathpisave_for_backwardconfig)
ctx
query_maskkey_maskrH   rI   valuerU   rN   expectationcumulation_valuer   r   r"   forwardo   s   (0zYosoCumulation.forwardc                 C   s   t |}| j\}}}}}}| j}|d }	t||dd| }
t|
|	d | }t|
dd|	d | }t|dd|}d d |||d fS )NrN   r8   rO   r7   )r6   saved_tensorsrU   rB   rD   rQ   )rV   gradrW   rX   rZ   rH   rI   rY   rU   rN   weighted_exp
grad_querygrad_key
grad_valuer   r   r"   backward|   s   zYosoCumulation.backwardN__name__
__module____qualname__staticmethodr\   rc   r   r   r   r"   rM   n   s
    
rM   c                   @   rL   )YosoLSHCumulationc              
   C   sV  | d| dkrtd| d| dkrtd| d| dkr*td| d| dkr8td| d| dkrFtd| d| dkrTtd	t|||||g\}}}}}|j}|d
 }|d }	td|	 }
|d rt||||||	|d\}}n	t||||	\}}t||||||
|d}| ||||||| || _	|S )Nr   z6Query mask and Key mask differ in sizes in dimension 0z3Query mask and Query differ in sizes in dimension 0z1Query mask and Key differ in sizes in dimension 0z8Query mask and Value mask differ in sizes in dimension 0r   z,Key and Value differ in sizes in dimension 1r7   z,Query and Key differ in sizes in dimension 2rJ   rN   use_fast_hash)
r@   rA   r6   is_cudarF   r,   Z	fast_hashrK   rT   rU   )rV   rW   rX   rH   rI   rY   rU   use_cudarJ   rN   hashtable_capacityquery_hash_codekey_hash_coder[   r   r   r"   r\      s8   
zYosoLSHCumulation.forwardc                 C   sj  t |}| j\}}}}}}}| j}	|j}
|	d }td| }|	d rSt|||||||
d}t|||||||d | ||
d
}t|||||||d | ||
d
}nZdtt	||
ddtj  | }||d d d d d f  |d d d d d f  }t	||
dd| }t	||d | }t	|
dd|d | }t	|
dd|}d d |||d fS )NrN   r7   lsh_backwardr      r8   rO   )r6   r]   rU   rk   rF   r,   Zlsh_weighted_cumulationrB   rP   rD   rQ   rR   rS   )rV   r^   rW   rX   rn   ro   rH   rI   rY   rU   rl   rN   rm   rb   r`   ra   rZ   r_   r   r   r"   rc      sR   

(0zYosoLSHCumulation.backwardNrd   r   r   r   r"   ri      s
    
%ri   c                       s*   e Zd ZdZ fddZdddZ  ZS )YosoEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.c                    s   t    tj|j|j|jd| _t|jd |j| _	t|j
|j| _tj|j|jd| _t|j| _| jdt|jdd dd t|dd	| _| jd
tj| j tj| jjddd d S )N)padding_idxr7   Zepsposition_ids)r   r8   F)
persistentposition_embedding_typeabsolutetoken_type_idsdtyper=   )super__init__r   	Embedding
vocab_sizehidden_sizeZpad_token_idword_embeddingsZmax_position_embeddingsposition_embeddingsZtype_vocab_sizetoken_type_embeddings	LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutZregister_bufferrB   rC   expandgetattrrw   zerosru   r@   longr=   selfrU   	__class__r   r"   r}      s   

zYosoEmbeddings.__init__Nc                 C   s   |d ur	|  }n|  d d }|d }|d u r$| jd d d |f }|d u rNt| drC| jd d d |f }||d |}|}ntj|tj| jjd}|d u rW| 	|}| 
|}	||	 }
| jdkrn| |}|
|7 }
| |
}
| |
}
|
S )Nr8   r   ry   r   rz   rx   )r@   ru   hasattrry   r   rB   r   r   r=   r   r   rw   r   r   r   )r   	input_idsry   ru   inputs_embedsinput_shape
seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr   
embeddingsr   r   r   r"   r\      s,   







zYosoEmbeddings.forward)NNNNre   rf   rg   __doc__r}   r\   __classcell__r   r   r   r"   rr      s    rr   c                       0   e Zd Zd	 fdd	Zdd Zd
ddZ  ZS )YosoSelfAttentionNc              
      s  t    |j|j dkrt|dstd|j d|j dtd u}t rKt rK|sKzt	  W n t
yJ } ztd|  W Y d }~nd }~ww |j| _t|j|j | _| j| j | _t|j| j| _t|j| j| _t|j| j| _t|j| _|d ur|n|j| _|j| _|j| _|jd u| _|j| _|j| _|j| _| j| j| j| jd| _ |jd urtj!|j|j|jdf|jd	 dfd
|jd| _"d S d S )Nr   Zembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()zGCould not load the custom kernel for multi-scale deformable attention: )rN   rj   rJ   rp   r   r7   F)Zin_channelsZout_channelsZkernel_sizepaddingbiasgroups)#r|   r}   r   num_attention_headsr   rA   r,   r   r   r-   	ExceptionloggerwarningrF   attention_head_sizeall_head_sizer   LinearrH   rI   rY   r   Zattention_probs_dropout_probr   rw   use_expectationrN   Zconv_windowuse_convrj   rJ   rp   
lsh_configZConv2dconv)r   rU   rw   Zkernel_loadeder   r   r"   r}   $  sZ   



zYosoSelfAttention.__init__c                 C   s6   |  d d | j| jf }|j| }|ddddS )Nr8   r   r7   r   r
   )r@   r   r   viewpermute)r   layerZnew_layer_shaper   r   r"   transpose_for_scoresW  s   
z&YosoSelfAttention.transpose_for_scoresFc                 C   s@  |  |}| | |}| | |}| |}| jr.| ||d d d d d d f  }| \}	}
}}||	|
 ||}||	|
 ||}||	|
 ||}d|d  }|dj	|
dd|	|
 |
 }d}| js||k r|	|
 ||| f}tj|tj||jdgdd}tj|tj||jdgdd}tj|tj||jdgdd}| js| jrt||g\}}| jrt|||||| j}nt|||||| j}| js||k r|d d d d d |f }t|}||	|
||}| jr||7 }|dd	dd
 }| d d | jf }|j| }|r||f}|S |f}|S )N      ?g     @r   r>       r<   r8   r   r7   r
   rO   )rH   r   rI   rY   r   r   r@   rE   Z	unsqueezeZrepeat_interleaverF   r   rB   catr   r=   trainingr;   rM   applyr   ri   r   r1   r   r   )r   hidden_statesattention_maskoutput_attentionsZmixed_query_layerZ	key_layerZvalue_layerZquery_layerZconv_value_layer
batch_sizeZ	num_headsZseq_lenZhead_dimZgpu_warp_sizeZpad_sizeZcontext_layerZnew_context_layer_shapeoutputsr   r   r"   r\   \  sx   

"	
zYosoSelfAttention.forwardr.   NF)re   rf   rg   r}   r   r\   r   r   r   r   r"   r   #  s    3r   c                       8   e Zd Z fddZdejdejdejfddZ  ZS )YosoSelfOutputc                    sB   t    t|j|j| _tj|j|jd| _t|j	| _
d S Nrt   )r|   r}   r   r   r   denser   r   r   r   r   r   r   r   r"   r}        
zYosoSelfOutput.__init__r   input_tensorreturnc                 C   &   |  |}| |}| || }|S r.   r   r   r   r   r   r   r   r   r"   r\        

zYosoSelfOutput.forwardre   rf   rg   r}   rB   Tensorr\   r   r   r   r   r"   r         $r   c                       r   )YosoAttentionNc                    s.   t    t||d| _t|| _t | _d S )N)rw   )r|   r}   r   r   r   outputsetpruned_heads)r   rU   rw   r   r   r"   r}     s   

zYosoAttention.__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>   )r?   r   r   r   r   r   r   rH   rI   rY   r   r   r   union)r   headsindexr   r   r"   prune_heads  s   zYosoAttention.prune_headsFc                 C   s4   |  |||}| |d |}|f|dd   }|S )Nr   r   )r   r   )r   r   r   r   Zself_outputsattention_outputr   r   r   r"   r\     s   zYosoAttention.forwardr.   r   )re   rf   rg   r}   r   r\   r   r   r   r   r"   r     s    r   c                       2   e Zd Z fddZdejdejfddZ  ZS )YosoIntermediatec                    sD   t    t|j|j| _t|jt	rt
|j | _d S |j| _d S r.   )r|   r}   r   r   r   intermediate_sizer   r/   
hidden_actstrr   intermediate_act_fnr   r   r   r"   r}     s
   
zYosoIntermediate.__init__r   r   c                 C      |  |}| |}|S r.   )r   r   r   r   r   r   r"   r\        

zYosoIntermediate.forwardr   r   r   r   r"   r     s    r   c                       r   )
YosoOutputc                    sB   t    t|j|j| _tj|j|jd| _t	|j
| _d S r   )r|   r}   r   r   r   r   r   r   r   r   r   r   r   r   r   r"   r}     r   zYosoOutput.__init__r   r   r   c                 C   r   r.   r   r   r   r   r"   r\     r   zYosoOutput.forwardr   r   r   r   r"   r     r   r   c                       s.   e Zd Z fddZd	ddZdd Z  ZS )
	YosoLayerc                    sB   t    |j| _d| _t|| _|j| _t|| _t	|| _
d S Nr   )r|   r}   chunk_size_feed_forwardseq_len_dimr   	attentionZadd_cross_attentionr   intermediater   r   r   r   r   r"   r}     s   


zYosoLayer.__init__NFc                 C   sF   | j |||d}|d }|dd  }t| j| j| j|}|f| }|S )N)r   r   r   )r   r   feed_forward_chunkr   r   )r   r   r   r   Zself_attention_outputsr   r   layer_outputr   r   r"   r\     s   
zYosoLayer.forwardc                 C   s   |  |}| ||}|S r.   )r   r   )r   r   Zintermediate_outputr   r   r   r"   r     s   
zYosoLayer.feed_forward_chunkr   )re   rf   rg   r}   r\   r   r   r   r   r   r"   r     s    
	r   c                       s0   e Zd Z fddZ					dddZ  ZS )	YosoEncoderc                    s:   t     | _t fddt jD | _d| _d S )Nc                    s   g | ]}t  qS r   )r   )r   _rU   r   r"   r#     r$   z(YosoEncoder.__init__.<locals>.<listcomp>F)	r|   r}   rU   r   Z
ModuleListrangenum_hidden_layersr   gradient_checkpointingr   r   r   r"   r}     s   
 
zYosoEncoder.__init__NFTc                 C   s   |rdnd }|r
dnd }t | jD ].\}	}
|r||f }| jr,| jr,| |
j|||}n|
|||}|d }|r?||d f }q|rG||f }|sUtdd |||fD S t|||dS )Nr   r   r   c                 s   s    | ]	}|d ur|V  qd S r.   r   )r   vr   r   r"   	<genexpr>D  s    z&YosoEncoder.forward.<locals>.<genexpr>)last_hidden_stater   
attentions)	enumerater   r   r   Z_gradient_checkpointing_func__call__tupler   )r   r   r   	head_maskr   output_hidden_statesreturn_dictZall_hidden_statesZall_self_attentionsiZlayer_moduleZlayer_outputsr   r   r"   r\   "  s4   	

zYosoEncoder.forward)NNFFT)re   rf   rg   r}   r\   r   r   r   r   r"   r     s    	r   c                       r   )YosoPredictionHeadTransformc                    sV   t    t|j|j| _t|jtrt	|j | _
n|j| _
tj|j|jd| _d S r   )r|   r}   r   r   r   r   r/   r   r   r   transform_act_fnr   r   r   r   r   r"   r}   N  s   
z$YosoPredictionHeadTransform.__init__r   r   c                 C   s"   |  |}| |}| |}|S r.   )r   r   r   r   r   r   r"   r\   W  s   


z#YosoPredictionHeadTransform.forwardr   r   r   r   r"   r   M  s    	r   c                       s,   e Zd Z fddZdd Zdd Z  ZS )YosoLMPredictionHeadc                    sL   t    t|| _tj|j|jdd| _t	t
|j| _| j| j_d S )NF)r   )r|   r}   r   	transformr   r   r   r   decoder	ParameterrB   r   r   r   r   r   r"   r}   `  s
   

zYosoLMPredictionHead.__init__c                 C   s   | j | j_ d S r.   )r   r   r   r   r   r"   _tie_weightsm  s   z!YosoLMPredictionHead._tie_weightsc                 C   r   r.   )r   r   r   r   r   r"   r\   p  r   zYosoLMPredictionHead.forward)re   rf   rg   r}   r   r\   r   r   r   r   r"   r   _  s    r   c                       r   )YosoOnlyMLMHeadc                    s   t    t|| _d S r.   )r|   r}   r   predictionsr   r   r   r"   r}   x  s   
zYosoOnlyMLMHead.__init__sequence_outputr   c                 C   s   |  |}|S r.   )r   )r   r  prediction_scoresr   r   r"   r\   |  s   
zYosoOnlyMLMHead.forwardr   r   r   r   r"   r   w  s    r   c                   @   s    e Zd ZeZdZdZdd ZdS )YosoPreTrainedModelr   Tc                 C   s   t |tjr |jjjd| jjd |jdur|jj	  dS dS t |tj
rC|jjjd| jjd |jdurA|jj|j 	  dS dS t |tjrX|jj	  |jjd dS dS )zInitialize the weightsg        )meanZstdNr   )r/   r   r   weightdataZnormal_rU   Zinitializer_ranger   Zzero_r~   rs   r   Zfill_)r   moduler   r   r"   _init_weights  s   

z!YosoPreTrainedModel._init_weightsN)re   rf   rg   r   Zconfig_classZbase_model_prefixZsupports_gradient_checkpointingr  r   r   r   r"   r    s
    r  c                       s   e Zd Z fddZdd Zdd Zdd Ze																		dd
ee	j
 dee	j
 dee	j
 dee	j
 dee	j
 dee	j
 dee dee dee deeef fddZ  ZS )	YosoModelc                    s2   t  | || _t|| _t|| _|   d S r.   )r|   r}   rU   rr   r   r   encoder	post_initr   r   r   r"   r}     s
   

zYosoModel.__init__c                 C   s   | j jS r.   r   r   r   r   r   r"   get_input_embeddings  s   zYosoModel.get_input_embeddingsc                 C   s   || j _d S r.   r  )r   rY   r   r   r"   set_input_embeddings  s   zYosoModel.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   )r   Zheads_to_pruner   r   r   r   r"   _prune_heads  s   zYosoModel._prune_headsNr   r   ry   ru   r   r   r   r   r   r   c
                 C   s  |d ur|n| j j}|d ur|n| j j}|	d ur|	n| j j}	|d ur*|d ur*td|d ur9| || | }
n|d urF| d d }
ntd|
\}}|d urU|jn|j}|d u retj	||f|d}|d u rt
| jdr| jjd d d |f }|||}|}n	tj|
tj|d}| || j j}| j||||d}| j||||||	d}|d	 }|	s|f|d
d   S t||j|j|jdS )NzDYou cannot specify both input_ids and inputs_embeds at the same timer8   z5You have to specify either input_ids or inputs_embedsr<   ry   rz   )r   ru   ry   r   )r   r   r   r   r   r   r   )r   r   r   cross_attentions)rU   r   r   use_return_dictrA   Z%warn_if_padding_and_no_attention_maskr@   r=   rB   Zonesr   r   ry   r   r   r   Zget_head_maskr   r
  r   r   r   r  )r   r   r   ry   ru   r   r   r   r   r   r   r   r   r=   r   r   Zembedding_outputZencoder_outputsr  r   r   r"   r\     s\   
zYosoModel.forward)	NNNNNNNNN)re   rf   rg   r}   r  r  r  r   r   rB   r   boolr   r   r   r\   r   r   r   r   r"   r	    sH    
	

r	  c                       s   e Zd ZddgZ fddZdd Zdd Ze																				dd
ee	j
 dee	j
 dee	j
 dee	j
 dee	j
 dee	j
 dee	j
 dee dee dee deeef fddZ  ZS )YosoForMaskedLMzcls.predictions.decoder.weightzcls.predictions.decoder.biasc                    s,   t  | t|| _t|| _|   d S r.   )r|   r}   r	  r   r   clsr  r   r   r   r"   r}     s   

zYosoForMaskedLM.__init__c                 C   s
   | j jjS r.   )r  r   r   r   r   r   r"   get_output_embeddings  s   
z%YosoForMaskedLM.get_output_embeddingsc                 C   s   || j j_|j| j j_d S r.   )r  r   r   r   )r   Znew_embeddingsr   r   r"   set_output_embeddings  s   
z%YosoForMaskedLM.set_output_embeddingsNr   r   ry   ru   r   r   labelsr   r   r   r   c                 C   s   |
dur|
n| j j}
| j||||||||	|
d	}|d }| |}d}|dur8t }||d| j j|d}|
sN|f|dd  }|durL|f| S |S t|||j|j	dS )a  
        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]`.
        Nr   ry   ru   r   r   r   r   r   r   r8   r   losslogitsr   r   )
rU   r  r   r  r   r   r   r   r   r   )r   r   r   ry   ru   r   r   r  r   r   r   r   r  r  Zmasked_lm_lossloss_fctr   r   r   r"   r\     s6   
zYosoForMaskedLM.forward
NNNNNNNNNN)re   rf   rg   Z_tied_weights_keysr}   r  r  r   r   rB   r   r  r   r   r   r\   r   r   r   r   r"   r    sN    		

r  c                       s(   e Zd ZdZ fddZdd Z  ZS )YosoClassificationHeadz-Head for sentence-level classification tasks.c                    sF   t    t|j|j| _t|j| _t|j|j	| _
|| _d S r.   )r|   r}   r   r   r   r   r   r   r   
num_labelsout_projrU   r   r   r   r"   r}   L  s
   

zYosoClassificationHead.__init__c                 K   sR   |d d dd d f }|  |}| |}t| jj |}|  |}| |}|S )Nr   )r   r   r   rU   r   r!  )r   featureskwargsxr   r   r"   r\   T  s   



zYosoClassificationHead.forwardr   r   r   r   r"   r  I  s    r  z
    YOSO Model transformer with a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks.
    )Zcustom_introc                          e Zd Z fddZe										ddeej deej deej deej deej d	eej d
eej dee dee dee de	e
ef fddZ  ZS )YosoForSequenceClassificationc                    s4   t  | |j| _t|| _t|| _|   d S r.   )r|   r}   r   r	  r   r  
classifierr  r   r   r   r"   r}   e  s
   

z&YosoForSequenceClassification.__init__Nr   r   ry   ru   r   r   r  r   r   r   r   c                 C   sh  |
dur|
n| j j}
| j||||||||	|
d	}|d }| |}d}|dur| j jdu rQ| jdkr7d| j _n| jdkrM|jtjksH|jtj	krMd| j _nd| j _| j jdkrot
 }| jdkri|| | }n+|||}n%| j jdkrt }||d| j|d}n| j jdkrt }|||}|
s|f|dd  }|dur|f| S |S t|||j|jd	S )
a  
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        Nr  r   r   Z
regressionZsingle_label_classificationZmulti_label_classificationr8   r  )rU   r  r   r'  Zproblem_typer   r{   rB   r   rF   r	   squeezer   r   r   r   r   r   )r   r   r   ry   ru   r   r   r  r   r   r   r   r  r  r  r  r   r   r   r"   r\   n  sT   


"


z%YosoForSequenceClassification.forwardr  )re   rf   rg   r}   r   r   rB   r   r  r   r   r   r\   r   r   r   r   r"   r&  ^  sH    		

r&  c                       r%  )YosoForMultipleChoicec                    sD   t  | t|| _t|j|j| _t|jd| _| 	  d S r   )
r|   r}   r	  r   r   r   r   pre_classifierr'  r  r   r   r   r"   r}     s
   
zYosoForMultipleChoice.__init__Nr   r   ry   ru   r   r   r  r   r   r   r   c                 C   s  |
dur|
n| j j}
|dur|jd n|jd }|dur%|d|dnd}|dur4|d|dnd}|durC|d|dnd}|durR|d|dnd}|dure|d|d|dnd}| j||||||||	|
d	}|d }|dddf }| |}t |}| 	|}|d|}d}|durt
 }|||}|
s|f|dd  }|dur|f| S |S t|||j|jdS )a[  
        input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`):
            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)
        token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *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)
        position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *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)
        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, 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.
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
            num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
            `input_ids` above)
        Nr   r8   rO   r  r   r  )rU   r  shaper   r@   r   r*  r   ZReLUr'  r   r   r   r   )r   r   r   ry   ru   r   r   r  r   r   r   Znum_choicesr   Zhidden_stateZpooled_outputr  Zreshaped_logitsr  r  r   r   r   r"   r\     sP   ,


zYosoForMultipleChoice.forwardr  )re   rf   rg   r}   r   r   rB   r   r  r   r   r   r\   r   r   r   r   r"   r)    sH    
	

r)  c                       r%  )YosoForTokenClassificationc                    sJ   t  | |j| _t|| _t|j| _t	|j
|j| _|   d S r.   )r|   r}   r   r	  r   r   r   r   r   r   r   r'  r  r   r   r   r"   r}   !  s   
z#YosoForTokenClassification.__init__Nr   r   ry   ru   r   r   r  r   r   r   r   c                 C   s
  |
dur|
n| j j}
| j||||||||	|
d	}|d }| |}| |}d}|duret }|durX|ddk}|d| j}t	||dt
|j|}|||}n||d| j|d}|
s{|f|dd  }|dury|f| S |S t|||j|jdS )z
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
        Nr  r   r8   r   r  )rU   r  r   r   r'  r   r   r   rB   wherer5   ignore_indexZtype_asr   r   r   )r   r   r   ry   ru   r   r   r  r   r   r   r   r  r  r  r  Zactive_lossZactive_logitsZactive_labelsr   r   r   r"   r\   ,  sF   

z"YosoForTokenClassification.forwardr  )re   rf   rg   r}   r   r   rB   r   r  r   r   r   r\   r   r   r   r   r"   r,    sH    	

r,  c                       s   e Zd Z fddZe											ddeej deej deej deej deej d	eej d
eej deej dee dee dee de	e
ef fddZ  ZS )YosoForQuestionAnsweringc                    sB   t  | d|_|j| _t|| _t|j|j| _| 	  d S )Nr7   )
r|   r}   r   r	  r   r   r   r   
qa_outputsr  r   r   r   r"   r}   m  s   
z!YosoForQuestionAnswering.__init__Nr   r   ry   ru   r   r   start_positionsend_positionsr   r   r   r   c                 C   s@  |d ur|n| j j}| j|||||||	|
|d	}|d }| |}|jddd\}}|d}|d}d }|d ur~|d ur~t| dkrK|d}t| dkrX|d}|d}|d|}|d|}t	|d}|||}|||}|| d }|s||f|dd   }|d ur|f| S |S t
||||j|jdS )	Nr  r   r   r8   r>   )r.  r7   )r  start_logits
end_logitsr   r   )rU   r  r   r0  splitr(  r?   r@   clampr   r   r   r   )r   r   r   ry   ru   r   r   r1  r2  r   r   r   r   r  r  r3  r4  Z
total_lossZignored_indexr  Z
start_lossZend_lossr   r   r   r"   r\   y  sP   








z YosoForQuestionAnswering.forward)NNNNNNNNNNN)re   rf   rg   r}   r   r   rB   r   r  r   r   r   r\   r   r   r   r   r"   r/  k  sN    	

r/  )r  r)  r/  r&  r,  r   r	  r  )Fr   rR   pathlibr   typingr   r   r   rB   Ztorch.utils.checkpointr   Ztorch.nnr   r   r	   Zactivationsr   Zmodeling_outputsr   r   r   r   r   r   Zmodeling_utilsr   Zpytorch_utilsr   r   r   utilsr   r   r   r   Zconfiguration_yosor   Z
get_loggerre   r   r,   r-   r6   r;   rK   ZautogradFunctionrM   ri   Modulerr   r   r   r   r   r   r   r   r   r   r   r  r	  r  r  r&  r)  r,  r/  __all__r   r   r   r"   <module>   sh    

Z< !2
fIQiKO