^j|dZddlmZddlmZddlZddlmZddlmZddl m Z dd l m Z dd lmZdd lmZdd lmZmZmZdd lmZmZddlmZddlmZddlmZmZm Z m!Z!m"Z"ddl#m$Z$ddl%m&Z&ddl'm(Z(e!jRe*Z+GddejXZ- d?dejXdej\dej\dej\dej\dzde/de/fdZ0Gdd ejXZ1Gd!d"ejXZ2Gd#d$ejXZ3Gd%d&ejXZ4Gd'd(ejXZ5Gd)d*eZ6Gd+d,ejXZ7eGd-d.eZ8eGd/d0e8Z9Gd1d2ejXZ:Gd3d4ejXZ;eGd5d6e8Z<ed78eGd9d:eZ=ed;8Gd<d=e8Z>gd>Z?y)@zPyTorch Splinter model.)Callable) dataclassN)nn)CrossEntropyLoss)initialization)ACT2FN)create_bidirectional_mask)GradientCheckpointingLayer)BaseModelOutput ModelOutputQuestionAnsweringModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)apply_chunking_to_forward)TransformersKwargsauto_docstringcan_return_tupleloggingtorch_compilable_check)merge_with_config_defaults)capture_outputs)SplinterConfigc eZdZdZfdZ d dej dzdej dzdej dzdejdzdef d Z xZ S) SplinterEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.c|t|tj|j|j |j |_tj|j|j |_ tj|j|j |_ tj|j |j|_ tj|j|_|j#dt%j&|jj)ddy)N) padding_idxeps position_idsrF) persistent)super__init__r Embedding vocab_size hidden_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutregister_buffertorcharangeexpandselfconfig __class__s y/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/transformers/models/splinter/modeling_splinter.pyr'zSplinterEmbeddings.__init__+s !||F,=,=v?Q?Q_e_r_rs#%<<0N0NPVPbPb#c %'\\&2H2H&J\J\%]"f&8&8f>S>STzz&"<"<=   ELL)G)GHOOPWXej  N input_idstoken_type_idsr" inputs_embedsreturnc||j}n|jdd}|d}||jddd|f}|:tj|tj|jj }||j |}|j|}||z}|j|} || z }|j|}|j|}|S)Nr$rdtypedevice) sizer"r7zeroslongrGr,r0r.r1r5) r;r@rAr"rB input_shape seq_lengthr0 embeddingsr.s r>forwardzSplinterEmbeddings.forward9s  #..*K',,.s3K ^  ,,Q  ^ J"%:: "66|D)) ^^J/ \\*- r?)NNNN) __name__ __module__ __qualname____doc__r'r7 LongTensor FloatTensortuplerN __classcell__r=s@r>rr(szQ  .2260426 ##d*((4/&&-  ((4/   r?rmodulequerykeyvalueattention_maskscalingr5c tj||jdd|z}|||z}tjj |dtj j|j}tjj|||j}tj||} | jddj} | |fS)Nrr$)dimrF)ptrainingr) r7matmul transposer functionalsoftmaxfloat32torFr5rb contiguous) rXrYrZr[r\r]r5kwargs attn_weights attn_outputs r>eager_attention_forwardrm[s<<s}}Q':;gEL!#n4 ==((2U]](SVVW\WbWbcL==((6??([L,,|U3K''1-88:K  $$r?c eZdZfdZ ddej dej dzdeede ej ej dzffdZ xZ S) SplinterSelfAttentionc$t||j|jzdk7r2t |ds&t d|jd|jd||_|j|_t|j|jz |_|j|jz|_ tj|j|j|_ tj|j|j|_ tj|j|j|_tj|j |_|j |_|jdz|_y)Nrembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()g)r&r'r*num_attention_headshasattr ValueErrorr<intattention_head_size all_head_sizerLinearrYrZr[r3attention_probs_dropout_probr5attention_dropoutr]r:s r>r'zSplinterSelfAttention.__init__ssC     : : :a ?PVXhHi#F$6$6#78 445Q8   #)#=#= #&v'9'9FrNzSplinterSelfAttention.forwards> $))#2. CCbC$*B*BC zz-055lCMMaQRS XXm,11,?II!QO zz-055lCMMaQRS (?(M(M KK , ,.E) %8      % $}}C$2H2HLL %  % ! \*k));;;;FFH L((r?N) rOrPrQr'r7TensorrTrrrUrNrVrWs@r>rororsd6048)||)))D0)+, ) u||U\\D00 1 )r?rocneZdZfdZdej dej dej fdZxZS)SplinterSelfOutputc(t|tj|j|j|_tj |j|j|_tj|j|_ yNr ) r&r'rryr*denser1r2r3r4r5r:s r>r'zSplinterSelfOutput.__init__s` YYv1163E3EF f&8&8f>S>STzz&"<"<= r?r| input_tensorrCcr|j|}|j|}|j||z}|Srrr5r1r;r|rs r>rNzSplinterSelfOutput.forward7 =1  ]3 }|'CD r?rOrPrQr'r7rrNrVrWs@r>rr1> U\\RWR^R^r?rc eZdZfdZ ddej dej dzdeedej fdZ xZ S) SplinterAttentioncbt|t||_t ||_yr)r&r'ror;routputr:s r>r'zSplinterAttention.__init__s& )&1 (0 r?Nr|r\rjrCc ^|}|j|fd|i|\}}|j||}|SNr\)r;r)r;r|r\rjresidual_s r>rNzSplinterAttention.forwardsK !$499  )   q  M8< r?r) rOrPrQr'r7rrTrrrNrVrWs@r>rrsQ148 || ))D0 +,   r?rcVeZdZfdZdej dej fdZxZS)SplinterIntermediatect|tj|j|j |_t|jtrt|j|_ y|j|_ yr) r&r'rryr*intermediate_sizer isinstance hidden_actstrr intermediate_act_fnr:s r>r'zSplinterIntermediate.__init__s] YYv1163K3KL f'' -'-f.?.?'@D $'-'8'8D $r?r|rCcJ|j|}|j|}|Sr)rr)r;r|s r>rNzSplinterIntermediate.forwards& =1 00? r?rrWs@r>rrs#9U\\ellr?rcneZdZfdZdej dej dej fdZxZS)SplinterOutputc(t|tj|j|j |_tj|j |j|_tj|j|_ yr) r&r'rryrr*rr1r2r3r4r5r:s r>r'zSplinterOutput.__init__s` YYv779K9KL f&8&8f>S>STzz&"<"<= r?r|rrCcr|j|}|j|}|j||z}|Srrrs r>rNzSplinterOutput.forwardrr?rrWs@r>rrrr?rc eZdZfdZ d dej dej dzdeedej fdZ dZ xZ S) SplinterLayerct||j|_d|_t ||_t ||_t||_ y)Nr) r&r'chunk_size_feed_forward seq_len_dimr attentionr intermediaterrr:s r>r'zSplinterLayer.__init__sI '-'E'E$*6208$V, r?Nr|r\rjrCc |j|fd|i|}t|j|j|j|}|Sr)rrfeed_forward_chunkrr)r;r|r\rjs r>rNzSplinterLayer.forwardsY '  )  2  # #T%A%A4CSCSUb r?cL|j|}|j||}|Sr)rr)r;attention_outputintermediate_output layer_outputs r>rz SplinterLayer.feed_forward_chunk s,"//0@A{{#68HI r?r) rOrPrQr'r7rrTrrrNrrVrWs@r>rrsV-48||))D0+,   $r?rc neZdZfdZ ddej dej dzdeede fdZ xZ S) SplinterEncoderct|||_tjt |j Dcgc] }t|c}|_d|_ ycc}w)NF) r&r'r<r ModuleListrangenum_hidden_layersrlayergradient_checkpointing)r;r<ir=s r>r'zSplinterEncoder.__init__sN  ]]5IaIaCb#caM&$9#cd &+#$dsA#Nr|r\rjrCc P|jD] }|||fi|}t|S)Nlast_hidden_state)rr )r;r|r\rj layer_modules r>rNzSplinterEncoder.forwardsC !JJ L(M +  r?r) rOrPrQr'r7rrTrrr rNrVrWs@r>rrsM,48 || ))D0 +,    r?rc<eZdZUeed<dZdZeedZ fdZ xZ S)SplinterPreTrainedModelr<splinterT)r| attentionsct||t|trZt j |j tj|j jdjdyy)Nr$r#) r& _init_weightsrrinitcopy_r"r7r8rr9)r;rXr=s r>rz%SplinterPreTrainedModel._init_weights3s[ f% f0 1 JJv**ELL9L9L9R9RSU9V,W,^,^_f,g h 2r?) rOrPrQr__annotations__base_model_prefixsupports_gradient_checkpointingrro_can_record_outputsrrVrWs@r>rr)s/ "&*#&+ iir?rc eZdZdZfdZdZdZeee dde jdzde jdzde jdzd e jdzd e jdzd e e d eezfd ZxZS) SplinterModela2 The model is an encoder (with only self-attention) following the architecture described in [Attention is all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. ct||||_t||_t ||_|jyr)r&r'r<rrMrencoder post_initr:s r>r'zSplinterModel.__init__As;   ,V4&v.  r?c.|jjSrrMr,)r;s r>get_input_embeddingsz"SplinterModel.get_input_embeddingsKs...r?c&||j_yrr)r;r[s r>set_input_embeddingsz"SplinterModel.set_input_embeddingsNs*/'r?Nr@r\rAr"rBrjrCc | | td|#|j|||j}n!||jdd}n td|\}} | |jn |j} |t j || f| }|&t j |tj| }|j||||} t|j| |}|j| fd |i|} | d } t| 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) NzDYou cannot specify both input_ids and inputs_embeds at the same timer$z5You have to specify either input_ids or inputs_embedsrGrE)r@r"rArB)r<rBr\r\rr) ru%warn_if_padding_and_no_attention_maskrHrGr7onesrIrJrMr r<rr )r;r@r\rAr"rBrjrK batch_sizerLrGembedding_outputencoder_outputssequence_outputs r>rNzSplinterModel.forwardQs96  ]%>cd d  "  6 6y. Q#..*K  &',,.s3KTU U!, J%.%:!! @T@T  !"ZZ*j)A6RN  !"[[EJJvVN??%)' + 3;;*)  '$,,  )   *!,-  r?)NNNNN)rOrPrQrRr'rrrrrr7rrrrUr rNrVrWs@r>rr9s /0 *..2.2,0-1 @ <<$&@  t+@  t+ @ llT) @ ||d* @ +,@   @  @ r?rcXeZdZdfd Zdej dej fdZxZS)SplinterFullyConnectedLayerct|||_||_t j |j|j|_t||_t j|j|_ yr) r&r' input_dim output_dimrryrr act_fnr1)r;rrrr=s r>r'z$SplinterFullyConnectedLayer.__init__sV "$YYt~~t? Z( doo6r?inputsrCcl|j|}|j|}|j|}|Sr)rrr1)r;rr|s r>rNz#SplinterFullyConnectedLayer.forwards2 6*  M2 }5 r?)gelurrWs@r>rrs#7ellu||r?rc(eZdZdZfdZdZxZS)QuestionAwareSpanSelectionHeadzf Implementation of Question-Aware Span Selection (QASS) head, described in Splinter's paper: ct|t|j|j|_t|j|j|_t|j|j|_t|j|j|_tj|j|jd|_ tj|j|jd|_ y)NF)bias) r&r'rr*query_start_transformquery_end_transformstart_transform end_transformrrystart_classifierend_classifierr:s r>r'z'QuestionAwareSpanSelectionHead.__init__s %@ASASU[UgUg%h"#>v?Q?QSYSeSe#f :6;M;MvOaOab89K9KVM_M_` " &*<*P>PW\ ] ii(:(:FrNz&QuestionAwareSpanSelectionHead.forwardsKKM 1c##B'..q!S9 V%@ 55mD11-@))&1 %%f---.>? ''1a0 ||M:> ++N; ##Aq!,\\-: Z''r?)rOrPrQrRr'rNrVrWs@r>rrs \(r?rcPeZdZfdZee ddejdzdejdzdejdzdejdzdejdzdejdzd ejdzd ejdzd e e d e e zfd Z xZS)SplinterForQuestionAnsweringct||t||_t ||_|j |_|jyrr&r'rrr splinter_qassquestion_token_idrr:s r>r'z%SplinterForQuestionAnswering.__init__C  %f- ;FC!'!9!9 r?Nr@r\rAr"rBstart_positions end_positionsquestion_positionsrjrCc  d} ||Dtjtj||jj d} nJtj |j dtj|j|j} | jd}d} |j|f||||d| } | d} |j| |\}}| r"|jd |jd }}|d|d |z tj|jj zz}|d |z tj|jj zz}d}||t#|j d kDr|jd}t#|j d kDr|jd}|j d }|j%d||j%d|t'| }|||}|||}||zd z }t)|||| j*| j, 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)`. FNr$)r`r)rFlayoutrGTr\rAr"rBr ignore_indexr_lossrrr|r)r7argmaxeqr rvrIrHrJrrGrrr squeezefinforFminlenclamp_rrr|r)r;r@r\rAr"rBr r rrjquestion_positions_were_none"question_position_for_each_exampleoutputsrrr total_loss ignored_indexloss_fct start_lossend_losss r>rNz$SplinterForQuestionAnswering.forwardsFD(-$  %$5:\\XXi)?)?@EEGR626;[[!&&q)MDXDXanauau62"D!M!Mb!Q +/ ($--  ))%'    "!*#'#5#5oGY#Z j ''3';';A'> @R@RST@U*L  %'1~+=\M_M_A`AdAd*ddL#q>'9U[[IYIY=Z=^=^&^^J  &=+D?'')*Q."1"9"9""==%%'(1, - 5 5b 9 (--a0M  " "1m 4  M 2']CH!,@J M:H$x/14J+%!!//))   r?NNNNNNNN)rOrPrQr'rrr7rrSrrrUrrNrVrWs@r>rrs*..2.2,0-137156:W <<$&W  t+W  t+ W llT) W ||d* W ))D0W ''$.W ",,t3W +,W  - -W W r?rzB Class for outputs of Splinter as a span selection model. ) custom_introceZdZUdZdZej dzed<dZej dzed<dZ ej dzed<dZ e ej dzed<dZ e ej dzed<y)SplinterForPreTrainingOutputa 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). Nrrrr|r) rOrPrQrRrr7rTrrrr|rUrr?r>r(r(7s|&*D%  d ")-1L%##d*1+/J!!D(/59M5**+d2926Je''(4/6r?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. ceZdZfdZee ddejdzdejdzdejdzdejdzdejdzdejdzd ejdzd ejdzd e e d e e zfd Z dejd ejfdZxZS)SplinterForPreTrainingct||t||_t ||_|j |_|jyrrr:s r>r'zSplinterForPreTraining.__init__Vr r?Nr@r\rAr"rBr r rrjrCc  || | td| | td||j|}|j|f||||d| } | d} | j\} } }|j | |\}}|jd}||j dj | || }|d|z tj|jjzz}|d|z tj|jjzz}d}|||jdtd| dz |jdtd| dz t|jj}||j!| |z| |j!| |z}||j!| |z| |j!| |z}||zdz }t#|||| j$| j& 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 inputs_embeds is usedrrrrr_r) TypeError_prepare_question_positionsrrHr rr9r7rrFrrmaxrr<r+rr(r|r)r;r@r\rAr"rBr r rrjrrrsequence_lengthr`rr num_questions attention_mask_for_each_questionr r"r#r$s r>rNzSplinterForPreTraining.forward`sDj  %/*E-Jcab b  'I,=]^ ^  '!%!A!A)!L $--  ))%'    "!*+:+?+?+A( OS#'#5#5oGY#Z j*//2  %/=/G/G/J/Q/QM?0 ,(1/O+OSXS^S^_k_q_qSrSvSv*vvL#q+K'Ku{{[e[k[kOlOpOp&ppJ  &=+D  " "1c!_q-@&A B  C?Q+>$? @ (T[[5M5MNH!!!*}"H%x/14J+%!!//))   r?cTtj||jjk(\}}tj|}t |j d|j dk(dtj||d}tj|j d|j|z }tj|j d|jf|jjtj|j}||||f<|S)Nrz?All samples in the batch must have at least one question token.left)siderrE)r7wherer<r bincountrrH searchsortedr8rGfullr0r+rJ)r;r@rowsflat_positionsr2 first_idxcolsrs r>r/z2SplinterForPreTraining._prepare_question_positionss${{9 8U8U+UVnt,    q !Y^^A%6 6 M &&tT? ||DIIaL= IJJ ^^A  1 1 3 4 KK $ $**##  !/ $*r?r%)rOrPrQr'rrr7rrSrrrUr(rNr/rVrWs@r>r+r+Ns"*..2.2,0-137156:n <<$&n  t+n  t+ n llT) n ||d* n ))D0n ''$.n ",,t3n +,n  - -n n `U\\ellr?r+)rr+rrr)r~)@rRcollections.abcr dataclassesrr7rtorch.nnrrr activationsr masking_utilsr modeling_layersr modeling_outputsr r rmodeling_utilsrrprocessing_utilsr pytorch_utilsrutilsrrrrr utils.genericrutils.output_capturingrconfiguration_splinterr get_loggerrOloggerModulerrfloatrmrorrrrrrrrrrrr(r+__all__r)r?r>rSs$! %&!69ZZF&6jj752   H %//t% II% <<% % << % LL4' %  %%.3)BII3)n   . 299  RYY .D bii 2 io i iZ +Z Z z"))$#(RYY#(Ld #:d d N  7;7  7"N4NNb r?