o
    i;                     @   s  d Z ddlZddlmZ ddlmZm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 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Z8G d*d+ d+ej1Z9G d,d- d-ej1Z:G d.d/ d/ej1Z;G d0d1 d1ej1Z<eG d2d3 d3eZ=eG d4d5 d5e=Z>eG d6d7 d7e=Z?G d8d9 d9ej1Z@ed:d;G d<d= d=e=ZAeG d>d? d?e=ZBeG d@dA dAe=ZCeG dBdC dCe=ZDg dDZEdS )EzPyTorch YOSO model.    N)Path)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss   )ACT2FN)GradientCheckpointingLayer)"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 )Nkernelsyosoc                    s   g | ]} | qS  r   ).0file
src_folderr   Z/home/ubuntu/.local/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)torch.utils.cpp_extensionr   r,   lsh_cumulation)r   r+   	src_filesr   r   r$   load_cuda_kernels4   s
   r1   c                 C   sJ   t | trg }| D ]}| s| }|| q	|S |  s#|  } | S N)
isinstancelistis_contiguous
contiguousappendinput_tensorsouttensorr   r   r$   to_contiguousC   s   
r<   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)r3   r4   r7   r   
functional	normalizer8   r   r   r$   rB   Q   s   
rB   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   r=   devicer   r>   r@   )lensize
ValueErrortorchrandnrD   arangematmulreshapeintsum)querykeynum_hashhash_lenrmat	raise_powquery_projectionkey_projectionquery_binary
key_binary
query_hashr   r   r$   hashing[   s   $$$r[   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   r>   )rI   acosrL   	transposemathpisave_for_backwardconfig)
ctx
query_maskkey_maskrP   rQ   valuere   r^   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 )Nr^   r>   r_   r=   )r<   saved_tensorsre   rI   rL   ra   )rf   gradrg   rh   rj   rP   rQ   ri   re   r^   weighted_exp
grad_querygrad_key
grad_valuer   r   r$   backward|   s   zYosoCumulation.backwardN__name__
__module____qualname__staticmethodrl   rs   r   r   r   r$   r]   n   s
    
r]   c                   @   r\   )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 1r=   z,Query and Key differ in sizes in dimension 2rR   r^   use_fast_hash)
rG   rH   r<   is_cudarN   r/   	fast_hashr[   rd   re   )rf   rg   rh   rP   rQ   ri   re   use_cudarR   r^   hashtable_capacityquery_hash_codekey_hash_coderk   r   r   r$   rl      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 )Nr^   r=   lsh_backwardr      r>   r_   )r<   rm   re   r{   rN   r/   lsh_weighted_cumulationrI   r`   rL   ra   rb   rc   )rf   rn   rg   rh   r   r   rP   rQ   ri   re   r}   r^   r~   rr   rp   rq   rj   ro   r   r   r$   rs      sR   

(0zYosoLSHCumulation.backwardNrt   r   r   r   r$   ry      s
    
%ry   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_idxr=   epsposition_ids)r   r>   F)
persistentposition_embedding_typeabsolutetoken_type_idsdtyperD   )super__init__r   	Embedding
vocab_sizehidden_sizepad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings	LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutregister_bufferrI   rK   expandgetattrr   zerosr   rG   longrD   selfre   	__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 )Nr>   r   r   r   r   r   )rG   r   hasattrr   r   rI   r   r   rD   r   r   r   r   r   r   )r   	input_idsr   r   inputs_embedsinput_shape
seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr   
embeddingsr   r   r   r$   rl      s,   







zYosoEmbeddings.forward)NNNNru   rv   rw   __doc__r   rl   __classcell__r   r   r   r$   r      s    r   c                       s(   e Zd Zd f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   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: )r^   rz   rR   r   r   r=   F)in_channelsout_channelskernel_sizepaddingbiasgroups)#r   r   r   num_attention_headsr   rH   r/   r   r   r1   	ExceptionloggerwarningrN   attention_head_sizeall_head_sizer   LinearrP   rQ   ri   r   attention_probs_dropout_probr   r   use_expectationr^   conv_windowuse_convrz   rR   r   
lsh_configConv2dconv)r   re   r   kernel_loadeder   r   r$   r   $  sZ   



zYosoSelfAttention.__init__Fc                 C   s  |j \}}}| ||d| j| jdd}| ||d| j| jdd}| ||d| j| jdd}	| jrM| 	|	|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 )Nr>   r   r=         ?g     @rE       rC   r   r	   r_   )shaperP   viewr   r   ra   rQ   ri   r   r   rG   rM   	unsqueezerepeat_interleaverN   r   rI   catr   rD   trainingrB   r]   applyr   ry   permuter6   r   )r   hidden_statesattention_maskoutput_attentions
batch_sizer   _query_layer	key_layervalue_layerconv_value_layer	num_headsseq_lenhead_dimgpu_warp_sizepad_sizecontext_layernew_context_layer_shapeoutputsr   r   r$   rl   W  s   "	
zYosoSelfAttention.forwardr2   NFru   rv   rw   r   rl   r   r   r   r   r$   r   #  s    3r   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 Nr   )r   r   r   r   r   denser   r   r   r   r   r   r   r   r$   r        
zYosoSelfOutput.__init__r   input_tensorreturnc                 C   &   |  |}| |}| || }|S r2   r   r   r   r   r   r   r   r   r$   rl        

zYosoSelfOutput.forwardru   rv   rw   r   rI   Tensorrl   r   r   r   r   r$   r         $r   c                       s0   e Zd Zd	 fdd	Zdd Zd
ddZ  ZS )YosoAttentionNc                    s.   t    t||d| _t|| _t | _d S )N)r   )r   r   r   r   r   outputsetpruned_heads)r   re   r   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   rE   )rF   r   r   r   r   r   r   rP   rQ   ri   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   self_outputsattention_outputr   r   r   r$   rl     s   zYosoAttention.forwardr2   r   )ru   rv   rw   r   r   rl   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 r2   )r   r   r   r   r   intermediate_sizer   r3   
hidden_actstrr
   intermediate_act_fnr   r   r   r$   r     s
   
zYosoIntermediate.__init__r   r   c                 C      |  |}| |}|S r2   )r   r  r   r   r   r   r$   rl        

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   r2   r   r   r   r   r$   rl     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   	attention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   self_attention_outputsr   r   layer_outputr   r   r$   rl     s   
zYosoLayer.forwardc                 C   s   |  |}| ||}|S r2   )r  r   )r   r   intermediate_outputr  r   r   r$   r    s   
zYosoLayer.feed_forward_chunkr   )ru   rv   rw   r   rl   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    r   re   r   r$   r%   %  r&   z(YosoEncoder.__init__.<locals>.<listcomp>F)	r   r   re   r   
ModuleListrangenum_hidden_layerslayer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 }|
|||}|d }|r/||d f }q|r7||f }|sEtdd |||fD S t|||dS )Nr   r   r   c                 s   s    | ]	}|d ur|V  qd S r2   r   )r    vr   r   r$   	<genexpr>B  s    z&YosoEncoder.forward.<locals>.<genexpr>)last_hidden_stater   
attentions)	enumerater  tupler   )r   r   r   	head_maskr   output_hidden_statesreturn_dictall_hidden_statesall_self_attentionsilayer_modulelayer_outputsr   r   r$   rl   (  s&   	

zYosoEncoder.forward)NNFFTr   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   r3   r  r  r
   transform_act_fnr   r   r   r   r   r$   r   L  s   
z$YosoPredictionHeadTransform.__init__r   r   c                 C   s"   |  |}| |}| |}|S r2   )r   r,  r   r  r   r   r$   rl   U  s   


z#YosoPredictionHeadTransform.forwardr   r   r   r   r$   r+  K  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	ParameterrI   r   r   r   r   r   r$   r   ^  s
   

zYosoLMPredictionHead.__init__c                 C   s   | j | j_ d S r2   )r   r/  r   r   r   r$   _tie_weightsk  s   z!YosoLMPredictionHead._tie_weightsc                 C   r  r2   )r.  r/  r  r   r   r$   rl   n  r	  zYosoLMPredictionHead.forward)ru   rv   rw   r   r2  rl   r   r   r   r   r$   r-  ]  s    r-  c                       r  )YosoOnlyMLMHeadc                    s   t    t|| _d S r2   )r   r   r-  predictionsr   r   r   r$   r   v  s   
zYosoOnlyMLMHead.__init__sequence_outputr   c                 C   s   |  |}|S r2   )r4  )r   r5  prediction_scoresr   r   r$   rl   z  s   
zYosoOnlyMLMHead.forwardr   r   r   r   r$   r3  u  s    r3  c                   @   s.   e Zd ZU eed< dZdZdejfddZ	dS )YosoPreTrainedModelre   r   Tmodulec                 C   s   | j j}t|tjr"|jjjd|d |jdur |jj	  dS dS t|tj
rC|jjjd|d |jdurA|jj|j 	  dS dS t|tjrX|jj	  |jjd dS t|tre|jj	  dS dS )zInitialize the weightsg        )meanstdNr   )re   initializer_ranger3   r   r   weightdatanormal_r   zero_r   r   r   fill_r-  )r   r8  r:  r   r   r$   _init_weights  s"   


z!YosoPreTrainedModel._init_weightsN)
ru   rv   rw   r   __annotations__base_model_prefixsupports_gradient_checkpointingr   ModulerA  r   r   r   r$   r7    s
   
 r7  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 r2   )r   r   re   r   r   r  encoder	post_initr   r   r   r$   r     s
   

zYosoModel.__init__c                 C   s   | j jS r2   r   r   r1  r   r   r$   get_input_embeddings  s   zYosoModel.get_input_embeddingsc                 C   s   || j _d S r2   rI  )r   ri   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)itemsrG  r  r  r   )r   heads_to_pruner  r   r   r   r$   _prune_heads  s   zYosoModel._prune_headsNr   r   r   r   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 timer>   z5You have to specify either input_ids or inputs_embedsrC   r   r   )r   r   r   r   )r   r#  r   r$  r%  r   r   )r  r   r   cross_attentions)re   r   r$  use_return_dictrH   %warn_if_padding_and_no_attention_maskrG   rD   rI   onesr   r   r   r   r   r   get_head_maskr  rG  r   r   r   rO  )r   r   r   r   r   r#  r   r   r$  r%  r   r   r   rD   r   r   embedding_outputencoder_outputsr5  r   r   r$   rl     s\   
zYosoModel.forward)	NNNNNNNNN)ru   rv   rw   r   rJ  rK  rN  r   r   rI   r   boolr   r"  r   rl   r   r   r   r   r$   rF    sH    
	

rF  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 r2   )r   r   rF  r   r3  clsrH  r   r   r   r$   r     s   

zYosoForMaskedLM.__init__c                 C   s
   | j jjS r2   )rX  r4  r/  r1  r   r   r$   get_output_embeddings  s   
z%YosoForMaskedLM.get_output_embeddingsc                 C   s   || j j_|j| j j_d S r2   )rX  r4  r/  r   )r   new_embeddingsr   r   r$   set_output_embeddings  s   
z%YosoForMaskedLM.set_output_embeddingsNr   r   r   r   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   r   r   r#  r   r   r$  r%  r   r>   r   losslogitsr   r   )
re   rP  r   rX  r   r   r   r   r   r   )r   r   r   r   r   r#  r   r\  r   r$  r%  r   r5  r6  masked_lm_lossloss_fctr   r   r   r$   rl     s6   
zYosoForMaskedLM.forward
NNNNNNNNNN)ru   rv   rw   _tied_weights_keysr   rY  r[  r   r   rI   r   rV  r   r"  r   rl   r   r   r   r   r$   rW     sN    		

rW  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 r2   )r   r   r   r   r   r   r   r   r   
num_labelsout_projre   r   r   r   r$   r   M  s
   

zYosoClassificationHead.__init__c                 K   sR   |d d dd d f }|  |}| |}t| jj |}|  |}| |}|S )Nr   )r   r   r
   re   r  rg  )r   featureskwargsxr   r   r$   rl   U  s   



zYosoClassificationHead.forwardr   r   r   r   r$   re  J  s    re  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.
    )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 r2   )r   r   rf  rF  r   re  
classifierrH  r   r   r   r$   r   f  s
   

z&YosoForSequenceClassification.__init__Nr   r   r   r   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   
regressionsingle_label_classificationmulti_label_classificationr>   r^  )re   rP  r   rn  problem_typerf  r   rI   r   rN   r   squeezer   r   r   r   r   r   )r   r   r   r   r   r#  r   r\  r   r$  r%  r   r5  r`  r_  rb  r   r   r   r$   rl   o  sT   


"


z%YosoForSequenceClassification.forwardrc  )ru   rv   rw   r   r   r   rI   r   rV  r   r"  r   rl   r   r   r   r   r$   rm  _  sH    		

rm  c                       rl  )YosoForMultipleChoicec                    sD   t  | t|| _t|j|j| _t|jd| _| 	  d S r  )
r   r   rF  r   r   r   r   pre_classifierrn  rH  r   r   r   r$   r     s
   
zYosoForMultipleChoice.__init__Nr   r   r   r   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   r>   r_   r]  r   r^  )re   rP  r   r   rG   r   ru  r   ReLUrn  r   r   r   r   )r   r   r   r   r   r#  r   r\  r   r$  r%  num_choicesr   hidden_statepooled_outputr`  reshaped_logitsr_  rb  r   r   r   r$   rl     sP   ,


zYosoForMultipleChoice.forwardrc  )ru   rv   rw   r   r   r   rI   r   rV  r   r"  r   rl   r   r   r   r   r$   rt    sH    
	

rt  c                       rl  )YosoForTokenClassificationc                    sJ   t  | |j| _t|| _t|j| _t	|j
|j| _|   d S r2   )r   r   rf  rF  r   r   r   r   r   r   r   rn  rH  r   r   r   r$   r   "  s   
z#YosoForTokenClassification.__init__Nr   r   r   r   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   r>   r   r^  )re   rP  r   r   rn  r   r   rf  rI   wherer;   ignore_indextype_asr   r   r   )r   r   r   r   r   r#  r   r\  r   r$  r%  r   r5  r`  r_  rb  active_lossactive_logitsactive_labelsr   r   r   r$   rl   -  sF   

z"YosoForTokenClassification.forwardrc  )ru   rv   rw   r   r   r   rI   r   rV  r   r"  r   rl   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 )Nr=   )
r   r   rf  rF  r   r   r   r   
qa_outputsrH  r   r   r   r$   r   n  s   
z!YosoForQuestionAnswering.__init__Nr   r   r   r   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   r>   rE   )r}  r=   )r_  start_logits
end_logitsr   r   )re   rP  r   r  splitrs  rF   rG   clampr   r   r   r   )r   r   r   r   r   r#  r   r  r  r   r$  r%  r   r5  r`  r  r  
total_lossignored_indexrb  
start_lossend_lossr   r   r   r$   rl   z  sP   








z YosoForQuestionAnswering.forward)NNNNNNNNNNN)ru   rv   rw   r   r   r   rI   r   rV  r   r"  r   rl   r   r   r   r   r$   r  l  sN    	

r  )rW  rt  r  rm  r{  r  rF  r7  )Fr   rb   pathlibr   typingr   r   rI   r   torch.nnr   r   r   activationsr
   modeling_layersr   modeling_outputsr   r   r   r   r   r   modeling_utilsr   pytorch_utilsr   r   r   utilsr   r   r   r   configuration_yosor   
get_loggerru   r   r/   r1   r<   rB   r[   autogradFunctionr]   ry   rE  r   r   r   r   r  r
  r  r  r+  r-  r3  r7  rF  rW  re  rm  rt  r{  r  __all__r   r   r   r$   <module>   sh    

Z< !*
fIQiKO