o
    ZhF                     @   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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 ddlmZ ee Z!G dd dej"Z#G dd dej"Z$G dd dej"Z%de$iZ&G dd dej"Z'G dd dej"Z(G dd dej"Z)G dd dej"Z*G dd dej"Z+eG d d! d!eZ,eG d"d# d#e,Z-G d$d% d%ej"Z.G d&d' d'ej"Z/eG d(d) d)e,Z0eG d*d+ d+eZ1ed,d-G d.d/ d/e,Z2g d0Z3dS )1zPyTorch Splinter model.    N)	dataclass)ListOptionalTupleUnion)nn)CrossEntropyLoss   )ACT2FN))BaseModelOutputWithPastAndCrossAttentionsModelOutputQuestionAnsweringModelOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringlogging   )SplinterConfigc                       sj   e Zd ZdZ fddZ					ddeej deej deej d	eej d
ee	 de
fddZ  ZS )SplinterEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.c                    s   t    tj|j|j|jd| _t|j|j| _	t|j
|j| _tj|j|jd| _t|j| _| jdt|jddd t|dd| _d S )	N)padding_idxZepsposition_ids)r   F)
persistentposition_embedding_typeabsolute)super__init__r   	EmbeddingZ
vocab_sizehidden_sizepad_token_idword_embeddingsmax_position_embeddingsposition_embeddingsZtype_vocab_sizetoken_type_embeddings	LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutZregister_buffertorcharangeexpandgetattrr   selfconfig	__class__ ]/var/www/auris/lib/python3.10/site-packages/transformers/models/splinter/modeling_splinter.pyr   +   s   
zSplinterEmbeddings.__init__Nr   	input_idstoken_type_idsr   inputs_embedspast_key_values_lengthreturnc                 C   s   |d ur	|  }n|  d d }|d }|d u r&| jd d ||| f }|d u r5tj|tj| jjd}|d u r>| |}| |}|| }	| jdkrU| 	|}
|	|
7 }	| 
|	}	| |	}	|	S )Nr   r   dtypedevicer   )sizer   r,   zeroslongr>   r#   r&   r   r%   r'   r+   )r1   r7   r8   r   r9   r:   input_shape
seq_lengthr&   
embeddingsr%   r5   r5   r6   forward<   s$   






zSplinterEmbeddings.forward)NNNNr   )__name__
__module____qualname____doc__r   r   r,   
LongTensorFloatTensorintr   rE   __classcell__r5   r5   r3   r6   r   (   s*    r   c                       s   e Zd Zd fdd	ZdejdejfddZ						dd	ejd
eej deej deej deej dee	e	ej   dee
 de	ej fddZ  ZS )SplinterSelfAttentionNc                    s   t    |j|j dkrt|dstd|j d|j d|j| _t|j|j | _| j| j | _t	
|j| j| _t	
|j| j| _t	
|j| j| _t	|j| _|p\t|dd| _| jdksh| jd	kry|j| _t	d
|j d | j| _|j| _d S )Nr   Zembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()r   r   relative_keyrelative_key_query   r   )r   r   r!   num_attention_headshasattr
ValueErrorrL   attention_head_sizeall_head_sizer   Linearquerykeyvaluer)   Zattention_probs_dropout_probr+   r/   r   r$   r    distance_embedding
is_decoderr1   r2   r   r3   r5   r6   r   `   s*   

zSplinterSelfAttention.__init__xr;   c                 C   s6   |  d d | j| jf }||}|ddddS )Nr   r   rR   r   r	   )r?   rS   rV   viewpermute)r1   r_   Znew_x_shaper5   r5   r6   transpose_for_scoresz   s   
z*SplinterSelfAttention.transpose_for_scoresFhidden_statesattention_mask	head_maskencoder_hidden_statesencoder_attention_maskpast_key_valueoutput_attentionsc                 C   s  |  |}|d u}	|	r|d ur|d }
|d }|}nP|	r/| | |}
| | |}|}n;|d urZ| | |}
| | |}tj|d |
gdd}
tj|d |gdd}n| | |}
| | |}| |}|d u}| jrz|
|f}t||
dd}| j	dks| j	dkr	|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rtd||}|| }n| j	dkr	td||}td|
|}|| | }|t| j }|d ur|| }tjj|dd}| |}|d ur0|| }t||}|dddd }| d d | jf }||}|rX||fn|f}| jrd||f }|S )Nr   r   rR   dimr   rP   rQ   r<   )r=   zbhld,lrd->bhlrzbhrd,lrd->bhlrr	   ) rY   rb   rZ   r[   r,   catr]   matmulZ	transposer   shapeZtensorrA   r>   r`   r-   r\   r$   tor=   ZeinsummathsqrtrV   r   Z
functionalZsoftmaxr+   ra   
contiguousr?   rW   )r1   rc   rd   re   rf   rg   rh   ri   Zmixed_query_layerZis_cross_attentionZ	key_layerZvalue_layerZquery_layer	use_cacheZattention_scoresZquery_lengthZ
key_lengthZposition_ids_lZposition_ids_rZdistanceZpositional_embeddingZrelative_position_scoresZrelative_position_scores_queryZrelative_position_scores_keyZattention_probsZcontext_layerZnew_context_layer_shapeoutputsr5   r5   r6   rE      sn   









zSplinterSelfAttention.forwardNNNNNNF)rF   rG   rH   r   r,   Tensorrb   r   rK   r   boolrE   rM   r5   r5   r3   r6   rN   _   s4    	rN   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 )SplinterSelfOutputc                    sB   t    t|j|j| _tj|j|jd| _t|j	| _
d S Nr   )r   r   r   rX   r!   denser'   r(   r)   r*   r+   r0   r3   r5   r6   r         
zSplinterSelfOutput.__init__rc   input_tensorr;   c                 C   &   |  |}| |}| || }|S rv   r}   r+   r'   r1   rc   r   r5   r5   r6   rE         

zSplinterSelfOutput.forwardrF   rG   rH   r   r,   rx   rE   rM   r5   r5   r3   r6   r{          $r{   eagerc                       s   e Zd Zd fdd	Zdd Z						ddejdeej d	eej d
eej deej dee	e	ej   dee
 de	ej fddZ  ZS )SplinterAttentionNc                    s4   t    t|j ||d| _t|| _t | _d S )Nr   )	r   r   SPLINTER_SELF_ATTENTION_CLASSESZ_attn_implementationr1   r{   outputsetpruned_headsr^   r3   r5   r6   r      s   

zSplinterAttention.__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   rj   )lenr   r1   rS   rV   r   r   rY   rZ   r[   r   r}   rW   union)r1   headsindexr5   r5   r6   prune_heads  s   zSplinterAttention.prune_headsFrc   rd   re   rf   rg   rh   ri   r;   c              	   C   s<   |  |||||||}| |d |}	|	f|dd   }
|
S )Nr   r   )r1   r   )r1   rc   rd   re   rf   rg   rh   ri   Zself_outputsattention_outputru   r5   r5   r6   rE     s   
	zSplinterAttention.forwardrv   rw   )rF   rG   rH   r   r   r,   rx   r   rK   r   ry   rE   rM   r5   r5   r3   r6   r      s4    	r   c                       s2   e Zd Z fddZdejdejfddZ  ZS )SplinterIntermediatec                    sD   t    t|j|j| _t|jt	rt
|j | _d S |j| _d S rv   )r   r   r   rX   r!   intermediate_sizer}   
isinstance
hidden_actstrr
   intermediate_act_fnr0   r3   r5   r6   r   /  s
   
zSplinterIntermediate.__init__rc   r;   c                 C   s   |  |}| |}|S rv   )r}   r   )r1   rc   r5   r5   r6   rE   7  s   

zSplinterIntermediate.forwardr   r5   r5   r3   r6   r   .  s    r   c                       rz   )SplinterOutputc                    sB   t    t|j|j| _tj|j|jd| _t	|j
| _d S r|   )r   r   r   rX   r   r!   r}   r'   r(   r)   r*   r+   r0   r3   r5   r6   r   ?  r~   zSplinterOutput.__init__rc   r   r;   c                 C   r   rv   r   r   r5   r5   r6   rE   E  r   zSplinterOutput.forwardr   r5   r5   r3   r6   r   >  r   r   c                       s   e Zd Z fddZ						ddejdeej deej deej d	eej d
eeeej   dee	 deej fddZ
dd Z  ZS )SplinterLayerc                    sr   t    |j| _d| _t|| _|j| _|j| _| jr-| js&t|  dt|dd| _	t
|| _t|| _d S )Nr   z> should be used as a decoder model if cross attention is addedr   r   )r   r   chunk_size_feed_forwardseq_len_dimr   	attentionr]   add_cross_attentionrU   crossattentionr   intermediater   r   r0   r3   r5   r6   r   N  s   


zSplinterLayer.__init__NFrc   rd   re   rf   rg   rh   ri   r;   c              	   C   s  |d ur
|d d nd }| j |||||d}	|	d }
| jr(|	dd }|	d }n|	dd  }d }| jro|d urot| dsDtd|  d|d urN|d	d  nd }| |
||||||}|d }
||dd  }|d }|| }t| j| j| j|
}|f| }| jr||f }|S )
NrR   )ri   rh   r   r   r   r   z'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`rl   )	r   r]   rT   rU   r   r   feed_forward_chunkr   r   )r1   rc   rd   re   rf   rg   rh   ri   Zself_attn_past_key_valueZself_attention_outputsr   ru   Zpresent_key_valueZcross_attn_present_key_valueZcross_attn_past_key_valueZcross_attention_outputslayer_outputr5   r5   r6   rE   \  sP   


	

zSplinterLayer.forwardc                 C   s   |  |}| ||}|S rv   )r   r   )r1   r   Zintermediate_outputr   r5   r5   r6   r     s   
z SplinterLayer.feed_forward_chunkrw   )rF   rG   rH   r   r,   rx   r   rK   r   ry   rE   r   rM   r5   r5   r3   r6   r   M  s4    	
Ar   c                       s   e Zd Z fddZ									ddejdeej deej d	eej d
eej deeeej   dee	 dee	 dee	 dee	 de
eej ef fddZ  ZS )SplinterEncoderc                    s:   t     | _t fddt jD | _d| _d S )Nc                    s   g | ]}t  qS r5   )r   ).0_r2   r5   r6   
<listcomp>  s    z,SplinterEncoder.__init__.<locals>.<listcomp>F)	r   r   r2   r   Z
ModuleListrangenum_hidden_layerslayergradient_checkpointingr0   r3   r   r6   r     s   
 
zSplinterEncoder.__init__NFTrc   rd   re   rf   rg   past_key_valuesrt   ri   output_hidden_statesreturn_dictr;   c                 C   s^  |	rdnd }|r
dnd }|r| j jrdnd }| jr%| jr%|r%td d}|r)dnd }t| jD ]^\}}|	r;||f }|d urC|| nd }|d urM|| nd }| jrc| jrc| |j	|||||||}n
||||||||}|d }|rz||d f7 }|r||d f }| j jr||d f }q0|	r||f }|
st
dd	 |||||fD S t|||||d
S )Nr5   zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Fr   r   r   rR   c                 s   s    | ]	}|d ur|V  qd S rv   r5   )r   vr5   r5   r6   	<genexpr>  s    z*SplinterEncoder.forward.<locals>.<genexpr>Zlast_hidden_stater   rc   
attentionscross_attentions)r2   r   r   ZtrainingloggerZwarning_once	enumerater   Z_gradient_checkpointing_func__call__tupler   )r1   rc   rd   re   rf   rg   r   rt   ri   r   r   Zall_hidden_statesZall_self_attentionsZall_cross_attentionsZnext_decoder_cacheiZlayer_moduleZlayer_head_maskrh   Zlayer_outputsr5   r5   r6   rE     sz   


zSplinterEncoder.forward)	NNNNNNFFT)rF   rG   rH   r   r,   rx   r   rK   r   ry   r   r   rE   rM   r5   r5   r3   r6   r     sD    		
r   c                   @   s    e Zd ZeZdZdZdd ZdS )SplinterPreTrainedModelsplinterTc                 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stdNg      ?)r   r   rX   weightdataZnormal_r2   Zinitializer_rangebiasZzero_r    r   r'   Zfill_)r1   moduler5   r5   r6   _init_weights  s   

z%SplinterPreTrainedModel._init_weightsN)rF   rG   rH   r   Zconfig_classZbase_model_prefixZsupports_gradient_checkpointingr   r5   r5   r5   r6   r     s
    r   c                        s   e Zd 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
j de	e
j de	ee
j  de	e de	e de	e de	e deeef fddZ  ZS )SplinterModela*  
    The model is an encoder (with only self-attention) following the architecture described in [Attention is all you
    need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
    Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
    c                    s2   t  | || _t|| _t|| _|   d S rv   )r   r   r2   r   rD   r   encoder	post_initr0   r3   r5   r6   r      s
   

zSplinterModel.__init__c                 C   s   | j jS rv   rD   r#   )r1   r5   r5   r6   get_input_embeddings*  s   z"SplinterModel.get_input_embeddingsc                 C   s   || j _d S rv   r   )r1   r[   r5   r5   r6   set_input_embeddings-  s   z"SplinterModel.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   )r1   Zheads_to_pruner   r   r5   r5   r6   _prune_heads0  s   zSplinterModel._prune_headsNr7   rd   r8   r   re   r9   rf   rg   r   rt   ri   r   r   r;   c                 C   s  |dur|n| j j}|dur|n| j j}|dur|n| j j}| j jr-|
dur(|
n| j j}
nd}
|dur;|dur;td|durJ| || | }n|durW| dd }ntd|\}}|durf|j	n|j	}|	durv|	d d j
d nd}|du rtj||| f|d}|du rtj|tj|d	}| ||}| j jr|dur| \}}}||f}|du rtj||d}| |}nd}| || j j}| j|||||d
}| j||||||	|
|||d
}|d }|s|f|dd  S t||j|j|j|jdS )a  
        token_type_ids (`torch.LongTensor` of shape `batch_size, 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, 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)
        NFzDYou cannot specify both input_ids and inputs_embeds at the same timer   z5You have to specify either input_ids or inputs_embedsr   rR   )r>   r<   )r7   r   r8   r9   r:   )	rd   re   rf   rg   r   rt   ri   r   r   r   r   )r2   ri   r   use_return_dictr]   rt   rU   Z%warn_if_padding_and_no_attention_maskr?   r>   ro   r,   Zonesr@   rA   Zget_extended_attention_maskZinvert_attention_maskZget_head_maskr   rD   r   r   r   rc   r   r   )r1   r7   rd   r8   r   re   r9   rf   rg   r   rt   ri   r   r   rB   
batch_sizerC   r>   r:   Zextended_attention_maskZencoder_batch_sizeZencoder_sequence_lengthr   Zencoder_hidden_shapeZencoder_extended_attention_maskZembedding_outputZencoder_outputssequence_outputr5   r5   r6   rE   8  sx    
zSplinterModel.forward)NNNNNNNNNNNNN)rF   rG   rH   rI   r   r   r   r   r   r   r,   rx   r   rK   ry   r   r   r   rE   rM   r5   r5   r3   r6   r     sb    
	

r   c                       s4   e Zd Zd fdd	ZdejdejfddZ  ZS )	SplinterFullyConnectedLayergeluc                    sD   t    || _|| _t| j| j| _t| | _t	| j| _	d S rv   )
r   r   	input_dim
output_dimr   rX   r}   r
   act_fnr'   )r1   r   r   r   r3   r5   r6   r     s   

z$SplinterFullyConnectedLayer.__init__inputsr;   c                 C   s"   |  |}| |}| |}|S rv   )r}   r   r'   )r1   r   rc   r5   r5   r6   rE     s   


z#SplinterFullyConnectedLayer.forward)r   r   r5   r5   r3   r6   r     s    
r   c                       s(   e Zd ZdZ fddZdd Z  ZS )QuestionAwareSpanSelectionHeadzf
    Implementation of Question-Aware Span Selection (QASS) head, described in Splinter's paper:

    c                    sz   t    t|j|j| _t|j|j| _t|j|j| _t|j|j| _tj	|j|jdd| _
tj	|j|jdd| _d S )NF)r   )r   r   r   r!   query_start_transformquery_end_transformstart_transformend_transformr   rX   start_classifierend_classifierr0   r3   r5   r6   r     s   
z'QuestionAwareSpanSelectionHead.__init__c                 C   s   |  \}}}|ddd|}tj|d|d}| |}| |}| |}	| |}
| 	|}|	
ddd}	t||	}| |}|

ddd}
t||
}||fS )Nr   r   )rk   r   r   rR   )r?   	unsqueezerepeatr,   gatherr   r   r   r   r   ra   rn   r   )r1   r   	positionsr   rk   r   Zgathered_repsZquery_start_repsZquery_end_repsZ
start_repsZend_repsrc   start_logits
end_logitsr5   r5   r6   rE     s   





z&QuestionAwareSpanSelectionHead.forward)rF   rG   rH   rI   r   rE   rM   r5   r5   r3   r6   r     s    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j de
eef fddZ  ZS )SplinterForQuestionAnsweringc                    4   t  | t|| _t|| _|j| _|   d S rv   r   r   r   r   r   splinter_qassquestion_token_idr   r0   r3   r5   r6   r     
   

z%SplinterForQuestionAnswering.__init__Nr7   rd   r8   r   re   r9   start_positionsend_positionsri   r   r   question_positionsr;   c                 C   s  |dur|n| j j}d}|du r9|dur#tjt|| j dd}ntj|dtj	|j
|jd}|d}d}| j|||||||	|
|d	}|d }| ||\}}|r`|d	|d	}}|dur~|d	| t|jj  }|d	| t|jj  }d}|dur|durt| d	kr|d}t| d	kr|d}|d	}|d| |d| t|d
}|||}|||}|| d }|s||f|d	d  }|dur|f| S |S t||||j|jdS )a  
        token_type_ids (`torch.LongTensor` of shape `batch_size, 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, 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)
        question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):
            The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size,
            num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be
            the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size,
            sequence_length)`.
        NFr   rj   r   )r=   layoutr>   Trd   r8   r   re   r9   ri   r   r   r   Zignore_indexrR   lossr   r   rc   r   )r2   r   r,   Zargmaxeqr   rL   r@   r?   rA   r   r>   r   r   r   Zsqueezefinfor=   minr   clamp_r   r   rc   r   )r1   r7   rd   r8   r   re   r9   r   r   ri   r   r   r   Zquestion_positions_were_noneZ"question_position_for_each_exampleru   r   r   r   
total_lossZignored_indexloss_fct
start_lossend_lossr   r5   r5   r6   rE     sj   $






z$SplinterForQuestionAnswering.forwardNNNNNNNNNNNN)rF   rG   rH   r   r   r   r,   rx   rJ   ry   r   r   r   rE   rM   r5   r5   r3   r6   r     sT    
	

r   c                   @   st   e Zd ZU dZdZeej ed< dZ	eej ed< dZ
eej ed< dZeeej  ed< dZeeej  ed< dS )SplinterForPreTrainingOutputa  
    Class for outputs of Splinter as a span selection model.

    Args:
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when start and end positions are provided):
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        start_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`):
            Span-start scores (before SoftMax).
        end_logits (`torch.FloatTensor` of shape `(batch_size, num_questions, sequence_length)`):
            Span-end scores (before SoftMax).
        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.
    Nr   r   r   rc   r   )rF   rG   rH   rI   r   r   r,   rK   __annotations__r   r   rc   r   r   r5   r5   r5   r6   r   Z  s   
 r   z
    Splinter Model for the recurring span selection task as done during the pretraining. The difference to the QA task
    is that we do not have a question, but multiple question tokens that replace the occurrences of recurring spans
    instead.
    )Zcustom_introc                       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j de
eef fddZdejdejfddZ  ZS )SplinterForPreTrainingc                    r   rv   r   r0   r3   r5   r6   r     r   zSplinterForPreTraining.__init__Nr7   rd   r8   r   re   r9   r   r   ri   r   r   r   r;   c                 C   s  |dur|n| j j}|du r|dur|durtd|du r&|du r&td|du r/| |}| j|||||||	|
|d	}|d }| \}}}| ||\}}|d}|dur}|d|||}|d| t	
|jj  }|d| t	
|jj  }d}|dur|dur|dtd|d  |dtd|d  t| j jd}|||| |||| }|||| |||| }|| d }|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_questions, 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_questions, 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_questions, 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_questions, 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.
        start_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        question_positions (`torch.LongTensor` of shape `(batch_size, num_questions)`, *optional*):
            The positions of all question tokens. If given, start_logits and end_logits will be of shape `(batch_size,
            num_questions, sequence_length)`. If None, the first question token in each sequence in the batch will be
            the only one for which start_logits and end_logits are calculated and they will be of shape `(batch_size,
            sequence_length)`.
        NzCquestion_positions must be specified in order to calculate the lossz>question_positions must be specified when input_embeds is usedr   r   r   r   rR   r   )r2   r   	TypeError_prepare_question_positionsr   r?   r   r   r.   r,   r   r=   r   r   maxr   r"   r`   r   rc   r   )r1   r7   rd   r8   r   re   r9   r   r   ri   r   r   r   ru   r   r   Zsequence_lengthrk   r   r   num_questionsZ attention_mask_for_each_questionr   r   r   r   r   r5   r5   r6   rE     sh   7


zSplinterForPreTraining.forwardc                 C   sl   t || jjk\}}t |}t j|d| f| jjt j	|j
d}t dd |D }||||f< |S )Nr   r<   c                 S   s   g | ]}t |qS r5   )r,   r-   )r   nr5   r5   r6   r     s    zFSplinterForPreTraining._prepare_question_positions.<locals>.<listcomp>)r,   wherer2   r   Zbincountfullr?   r   r"   rA   r>   rm   )r1   r7   rowsZflat_positionsr   r   colsr5   r5   r6   r   	  s   
z2SplinterForPreTraining._prepare_question_positionsr   )rF   rG   rH   r   r   r   r,   rx   rJ   ry   r   r   r   rE   r   rM   r5   r5   r3   r6   r   z  sV    
	

|r   )r   r   r   r   r   )4rI   rq   dataclassesr   typingr   r   r   r   r,   Ztorch.utils.checkpointr   Ztorch.nnr   Zactivationsr
   Zmodeling_outputsr   r   r   Zmodeling_utilsr   Zpytorch_utilsr   r   r   utilsr   r   Zconfiguration_splinterr   Z
get_loggerrF   r   Moduler   rN   r{   r   r   r   r   r   r   r   r   r   r   r   r   r   __all__r5   r5   r5   r6   <module>   sV   
7 4W] &r 