o Zh@sdZddlZddlmZddlmZddlmZmZm Z ddl Z ddl Z ddl m Z ddl mZmZmZdd lmZmZerEdd lmZdd lmZdd lmZmZmZmZmZmZmZm Z m!Z!dd l"m#Z#ddl$m%Z%ddlm&Z&ddl'm(Z(e&)e*Z+ddZ,ddZ-ddZ.Gddde j/Z0Gddde j/Z1Gddde j/Z2Gddde j/Z3Gd d!d!e j/Z4Gd"d#d#e j/Z5Gd$d%d%e j/Z6Gd&d'd'e j/Z7Gd(d)d)e j/Z8Gd*d+d+e j/Z9Gd,d-d-e j/Z:Gd.d/d/e j/Z;Gd0d1d1e j/ZeGd6d7d7eZ?eGd8d9d9e>Z@ed:d;GdZAeGd>d?d?e>ZBed@d;GdAdBdBe>ZCedCd;GdDdEdEe>ZDeGdFdGdGe>ZEeGdHdIdIe>ZFeGdJdKdKe>ZGgdLZHdS)MzPyTorch FNet model.N) dataclass)partial)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)auto_docstringis_scipy_available)linalg)ACT2FN) BaseModelOutputBaseModelOutputWithPoolingMaskedLMOutput ModelOutputMultipleChoiceModelOutputNextSentencePredictorOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward)logging) FNetConfigcCs:|jd}|d|d|f}|tj}td|||S)z4Applies 2D matrix multiplication to 3D input arrays.rNzbij,jk,ni->bnk)shapetypetorch complex64Zeinsum)xmatrix_dim_onematrix_dim_two seq_lengthr&U/var/www/auris/lib/python3.10/site-packages/transformers/models/fnet/modeling_fnet.py_two_dim_matmul7s  r(cCs t|||SN)r()r"r#r$r&r&r'two_dim_matmul@ r*cCs4|}tt|jddD] }tjj||d}q |S)z Applies n-dimensional Fast Fourier Transform (FFT) to input array. Args: x: Input n-dimensional array. Returns: n-dimensional Fourier transform of input n-dimensional array. rN)axis)reversedrangendimr fft)r"outr,r&r&r'fftnEs r2cs*eZdZdZfddZdddZZS)FNetEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|j|_t|j|_|jdt|jddd|jdtj|jtjddddS) N) padding_idxZeps position_ids)rF) persistenttoken_type_idsdtype)super__init__r Embedding vocab_size hidden_sizeZ pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingsZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsLinear projectionDropouthidden_dropout_probdropoutregister_bufferr Zarangeexpandzerosr6sizelongselfconfig __class__r&r'r=Xs  zFNetEmbeddings.__init__Nc Cs|dur |}n|dd}|d}|dur$|jddd|f}|durNt|drC|jddd|f}||d|}|}n tj|tj|jjd}|durW| |}| |} || } | |} | | 7} | | } | | } || } | S)Nr7rr9rr;device)rOr6hasattrr9rMr rNrPrWrArDrCrErHrK) rR input_idsr9r6 inputs_embeds input_shaper%buffered_token_type_ids buffered_token_type_ids_expandedrD embeddingsrCr&r&r'forwardns,        zFNetEmbeddings.forward)NNNN)__name__ __module__ __qualname____doc__r=r_ __classcell__r&r&rTr'r3Us r3c,eZdZfddZddZddZZS)FNetBasicFourierTransformcst||dSr))r<r=_init_fourier_transformrQrTr&r'r= z"FNetBasicFourierTransform.__init__cCs|jsttjjdd|_dS|jdkrLtrB|dtj t |j tj d|dtj t |jtj dtt|j|jd|_dStdt|_dSt|_dS) N)rdimdft_mat_hiddenr: dft_mat_seq)r#r$zpSciPy is needed for DFT matrix calculation and is not found. Using TPU optimized fast fourier transform instead.)use_tpu_fourier_optimizationsrr r0r2fourier_transformrBr rLZtensorrZdftr@r!tpu_short_seq_lengthr*rnrmrwarningrQr&r&r'rgs$     z1FNetBasicFourierTransform._init_fourier_transformcCs||j}|fSr))rpreal)rR hidden_statesoutputsr&r&r'r_s z!FNetBasicFourierTransform.forward)r`rarbr=rgr_rdr&r&rTr'rfs rfc$eZdZfddZddZZS)FNetBasicOutputcs"ttj|j|jd|_dSNr5)r<r=rrEr@rFrQrTr&r'r=s zFNetBasicOutput.__init__cCs|||}|Sr))rErRrt input_tensorr&r&r'r_szFNetBasicOutput.forwardr`rarbr=r_rdr&r&rTr'rw rwcrv)FNetFourierTransformcs"tt||_t||_dSr))r<r=rfrRrwoutputrQrTr&r'r=s  zFNetFourierTransform.__init__cCs$||}||d|}|f}|SNr)rRr~)rRrtZ self_outputsfourier_outputrur&r&r'r_s zFNetFourierTransform.forwardr{r&r&rTr'r} r}c2eZdZfddZdejdejfddZZS)FNetIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr)) r<r=rrGr@intermediate_sizedense isinstance hidden_actstrrintermediate_act_fnrQrTr&r'r=s   zFNetIntermediate.__init__rtreturncC||}||}|Sr))rrrRrtr&r&r'r_  zFNetIntermediate.forwardr`rarbr=r Tensorr_rdr&r&rTr'rs rcs8eZdZfddZdejdejdejfddZZS) FNetOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSrx) r<r=rrGrr@rrErFrIrJrKrQrTr&r'r=s zFNetOutput.__init__rtrzrcCs&||}||}|||}|Sr))rrKrEryr&r&r'r_s  zFNetOutput.forwardrr&r&rTr'rs $rcre) FNetLayercs:t|j|_d|_t||_t||_t||_ dSNr) r<r=chunk_size_feed_forward seq_len_dimr}fourierr intermediaterr~rQrTr&r'r=s   zFNetLayer.__init__cCs0||}|d}t|j|j|j|}|f}|Sr)rrfeed_forward_chunkrr)rRrtZself_fourier_outputsr layer_outputrur&r&r'r_s zFNetLayer.forwardcCs||}|||}|Sr))rr~)rRrZintermediate_outputrr&r&r'rs  zFNetLayer.feed_forward_chunk)r`rarbr=r_rrdr&r&rTr'rs  rcs&eZdZfddZdddZZS) FNetEncodercs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr&)r).0_rSr&r' sz(FNetEncoder.__init__..F) r<r=rSrZ ModuleListr.Znum_hidden_layerslayergradient_checkpointingrQrTrr'r= s   zFNetEncoder.__init__FTcCs|rdnd}t|jD]!\}}|r||f}|jr$|jr$||j|}n||}|d}q |r4||f}|sAtdd||fDSt||dS)Nr&rcss|] }|dur|VqdSr)r&)rvr&r&r' "sz&FNetEncoder.forward..)last_hidden_statert) enumeraterrZtrainingZ_gradient_checkpointing_func__call__tupler)rRrtoutput_hidden_states return_dictZall_hidden_statesiZ layer_moduleZ layer_outputsr&r&r'r_s      zFNetEncoder.forward)FTr{r&r&rTr'r s rcr) FNetPoolercs*tt|j|j|_t|_dSr))r<r=rrGr@rZTanh activationrQrTr&r'r=)s zFNetPooler.__init__rtrcCs(|dddf}||}||}|Sr)rr)rRrtZfirst_token_tensor pooled_outputr&r&r'r_.s  zFNetPooler.forwardrr&r&rTr'r(s rcr)FNetPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dSrx) r<r=rrGr@rrrrrtransform_act_fnrErFrQrTr&r'r=9s  z$FNetPredictionHeadTransform.__init__rtrcCs"||}||}||}|Sr))rrrErr&r&r'r_Bs   z#FNetPredictionHeadTransform.forwardrr&r&rTr'r8s  rcs.eZdZfddZddZd ddZZS) FNetLMPredictionHeadcsHtt||_t|j|j|_t t |j|_ |j |j_ dSr)) r<r=r transformrrGr@r?decoder Parameterr rNbiasrQrTr&r'r=Js  zFNetLMPredictionHead.__init__cCrr))rrrr&r&r'r_UrzFNetLMPredictionHead.forwardrNcCs,|jjjjdkr|j|j_dS|jj|_dS)Nmeta)rrrWrrRr&r&r' _tie_weightsZsz!FNetLMPredictionHead._tie_weights)rN)r`rarbr=r_rrdr&r&rTr'rIs  rcrv)FNetOnlyMLMHeadcstt||_dSr))r<r=r predictionsrQrTr&r'r=drhzFNetOnlyMLMHead.__init__cC||}|Sr))r)rRsequence_outputprediction_scoresr&r&r'r_h zFNetOnlyMLMHead.forwardr{r&r&rTr'rcr|rcrv)FNetOnlyNSPHeadcstt|jd|_dSNri)r<r=rrGr@seq_relationshiprQrTr&r'r=os zFNetOnlyNSPHead.__init__cCrr))r)rRrseq_relationship_scorer&r&r'r_srzFNetOnlyNSPHead.forwardr{r&r&rTr'rnr|rcrv)FNetPreTrainingHeadscs(tt||_t|jd|_dSr)r<r=rrrrGr@rrQrTr&r'r=zs  zFNetPreTrainingHeads.__init__cCs||}||}||fSr))rr)rRrrrrr&r&r'r_s  zFNetPreTrainingHeads.forwardr{r&r&rTr'ryrrc@s eZdZeZdZdZddZdS)FNetPreTrainedModelfnetTcCst|tjr |jjjd|jjd|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|tj rX|jj |jj ddSdS)zInitialize the weightsg)meanZstdNg?)rrrGweightdataZnormal_rSZinitializer_rangerZzero_r>r4rEZfill_)rRmoduler&r&r' _init_weightss      z!FNetPreTrainedModel._init_weightsN)r`rarbrZ config_classZbase_model_prefixZsupports_gradient_checkpointingrr&r&r&r'rs  rc@s^eZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dS)FNetForPreTrainingOutputa Output type of [`FNetForPreTraining`]. Args: loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss. prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`): Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation 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 + 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 initial embedding outputs. Nlossprediction_logitsseq_relationship_logitsrt) r`rarbrcrrr FloatTensor__annotations__rrrtrr&r&r&r'rs rcseZdZdZdfdd ZddZddZe dd ee j d ee j d ee j d ee j dee dee de eeffddZZS) FNetModelz The model can behave as an encoder, following the architecture described in [FNet: Mixing Tokens with Fourier Transforms](https://arxiv.org/abs/2105.03824) by James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, Santiago Ontanon. TcsDt|||_t||_t||_|rt|nd|_| dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) r<r=rSr3r^rencoderrpooler post_init)rRrSZadd_pooling_layerrTr&r'r=s    zFNetModel.__init__cCs|jjSr)r^rArr&r&r'get_input_embeddingsszFNetModel.get_input_embeddingscCs ||j_dSr)r)rRvaluer&r&r'set_input_embeddingsr+zFNetModel.set_input_embeddingsNrYr9r6rZrrrcCsv|dur|n|jj}|dur|n|jj}|dur |dur td|dur-|}|\}} n|dur>|dd}|\}} ntd|jjrT| dkrT|jj| krTtd|dur[|jn|j} |durt|j dr}|j j ddd| f} | || } | }n t j |t j| d}|j ||||d} |j| ||d }|d }|jdur||nd}|s||f|d dSt|||jd S) NzDYou cannot specify both input_ids and inputs_embeds at the same timer7z5You have to specify either input_ids or inputs_embedsrlzThe `tpu_short_seq_length` in FNetConfig should be set equal to the sequence length being passed to the model when using TPU optimizations.r9rV)rYr6r9rZ)rrrr)r pooler_outputrt)rSruse_return_dict ValueErrorrOrorqrWrXr^r9rMr rNrPrrrrt)rRrYr9r6rZrrr[Z batch_sizer%rWr\r]Zembedding_outputZencoder_outputsrrr&r&r'r_s\      zFNetModel.forward)T)NNNNNN)r`rarbrcr=rrr rr Z LongTensorrboolrrrr_rdr&r&rTr'rs6 rz FNet Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next sentence prediction (classification)` head. )Z custom_introcseZdZddgZfddZddZddZe dd ee j d ee j d ee j d ee j dee j dee j dee dee de e effddZZS)FNetForPreTrainingcls.predictions.decoder.biascls.predictions.decoder.weightc,t|t||_t||_|dSr))r<r=rrrclsrrQrTr&r'r=&    zFNetForPreTraining.__init__cC |jjjSr)rrrrr&r&r'get_output_embeddings/ z(FNetForPreTraining.get_output_embeddingscC||jj_|j|jj_dSr)rrrrrRZnew_embeddingsr&r&r'set_output_embeddings2 z(FNetForPreTraining.set_output_embeddingsNrYr9r6rZlabelsnext_sentence_labelrrrc Cs|dur|n|jj}|j||||||d} | dd\} } || | \} } d}|durP|durPt}|| d|jj|d}|| dd|d}||}|sg| | f| dd}|dure|f|S|St|| | | jdS)aH 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]` next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see `input_ids` docstring) Indices should be in `[0, 1]`: - 0 indicates sequence B is a continuation of sequence A, - 1 indicates sequence B is a random sequence. Example: ```python >>> from transformers import AutoTokenizer, FNetForPreTraining >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base") >>> model = FNetForPreTraining.from_pretrained("google/fnet-base") >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") >>> outputs = model(**inputs) >>> prediction_logits = outputs.prediction_logits >>> seq_relationship_logits = outputs.seq_relationship_logits ```Nr9r6rZrrrir7)rrrrt) rSrrrr viewr?rrt)rRrYr9r6rZrrrrrurrrr total_lossloss_fctmasked_lm_lossnext_sentence_lossr~r&r&r'r_6s4% zFNetForPreTraining.forwardNNNNNNNN)r`rarb_tied_weights_keysr=rrr rr rrrrrr_rdr&r&rTr'rsB    rcseZdZddgZfddZddZddZe dd ee j d ee j d ee j d ee j dee j dee dee de e effddZZS)FNetForMaskedLMrrcrr))r<r=rrrrrrQrTr&r'r=rzFNetForMaskedLM.__init__cCrr)rrr&r&r'rrz%FNetForMaskedLM.get_output_embeddingscCrr)rrr&r&r'rrz%FNetForMaskedLM.set_output_embeddingsNrYr9r6rZrrrrcCs|dur|n|jj}|j||||||d}|d} || } d} |dur5t} | | d|jj|d} |sK| f|dd} | durI| f| S| St| | |jdS)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]`. Nrrr7rirlogitsrt) rSrrrr rr?rrt)rRrYr9r6rZrrrrurrrrr~r&r&r'r_s& zFNetForMaskedLM.forwardNNNNNNN)r`rarbrr=rrr rr rrrrrr_rdr&r&rTr'r|s<   rzT FNet Model with a `next sentence prediction (classification)` head on top. ceZdZfddZe       ddeejdeejdeejdeejdeejd eed eed e e e ffd d Z Z S)FNetForNextSentencePredictioncrr))r<r=rrrrrrQrTr&r'r=rz&FNetForNextSentencePrediction.__init__NrYr9r6rZrrrrcKsd|vrtdt|d}|dur|n|jj}|j||||||d} | d} || } d} |durBt} | | dd| d} |sX| f| dd}| durV| f|S|St | | | j dS) a labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see `input_ids` docstring). Indices should be in `[0, 1]`: - 0 indicates sequence B is a continuation of sequence A, - 1 indicates sequence B is a random sequence. Example: ```python >>> from transformers import AutoTokenizer, FNetForNextSentencePrediction >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base") >>> model = FNetForNextSentencePrediction.from_pretrained("google/fnet-base") >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") >>> outputs = model(**encoding, labels=torch.LongTensor([1])) >>> logits = outputs.logits >>> assert logits[0, 0] < logits[0, 1] # next sentence was random ```rzoThe `next_sentence_label` argument is deprecated and will be removed in a future version, use `labels` instead.Nrrr7rir) warningswarn FutureWarningpoprSrrrr rrrt)rRrYr9r6rZrrrkwargsrurZseq_relationship_scoresrrr~r&r&r'r_s:$  z%FNetForNextSentencePrediction.forwardr)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs6   rz FNet Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. cr)FNetForSequenceClassificationcJt||j|_t||_t|j|_t |j |j|_ | dSr) r<r= num_labelsrrrrIrJrKrGr@ classifierrrQrTr&r'r=s   z&FNetForSequenceClassification.__init__NrYr9r6rZrrrrcCsh|dur|n|jj}|j||||||d}|d} || } || } d} |dur|jjdurS|jdkr9d|j_n|jdkrO|jtj ksJ|jtj krOd|j_nd|j_|jjdkrqt } |jdkrk| | | } n+| | |} n%|jjdkrt } | | d|j|d} n|jjdkrt} | | |} |s| f|dd} | dur| f| S| St| | |jd 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). NrrZ regressionZsingle_label_classificationZmulti_label_classificationr7rir)rSrrrKrZ problem_typerr;r rPintr squeezer rrrrt)rRrYr9r6rZrrrrurrrrr~r&r&r'r_'sF     "       z%FNetForSequenceClassification.forwardr)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs6   rcr)FNetForMultipleChoicecs@t|t||_t|j|_t|j d|_ | dSr) r<r=rrrrIrJrKrGr@rrrQrTr&r'r=fs   zFNetForMultipleChoice.__init__NrYr9r6rZrrrrcCsF|dur|n|jj}|dur|jdn|jd}|dur%|d|dnd}|dur4|d|dnd}|durC|d|dnd}|durV|d|d|dnd}|j||||||d} | d} || } || } | d|} d} |durt}|| |} |s| f| dd}| dur| f|S|St | | | j dS)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) Nrr7rrir) rSrrrrOrrKrr rrt)rRrYr9r6rZrrrZ num_choicesrurrZreshaped_logitsrrr~r&r&r'r_ps:)    zFNetForMultipleChoice.forwardr)r`rarbr=r rr rrrrrr_rdr&r&rTr'rds6   rcr)FNetForTokenClassificationcrr)rrQrTr&r'r=s   z#FNetForTokenClassification.__init__NrYr9r6rZrrrrcCs|dur|n|jj}|j||||||d}|d} || } || } d} |dur9t} | | d|j|d} |sO| f|dd} | durM| f| S| St| | |j dS)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]`. Nrrr7rir) rSrrrKrr rrrrt)rRrYr9r6rZrrrrurrrrr~r&r&r'r_s(  z"FNetForTokenClassification.forwardr)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs6   rcseZdZfddZe        ddeejdeejdeejdeejdeejd eejd eed eed e e e ffd dZ Z S)FNetForQuestionAnsweringcs<t||j|_t||_t|j|j|_| dSr)) r<r=rrrrrGr@ qa_outputsrrQrTr&r'r=s   z!FNetForQuestionAnswering.__init__NrYr9r6rZstart_positions end_positionsrrrc Cs>|dur|n|jj}|j||||||d} | d} || } | jddd\} } | d} | d} d}|dur|durt|dkrL|d}t|dkrY|d}| d}| d|}| d|}t |d}|| |}|| |}||d}|s| | f| dd}|dur|f|S|St || | | j dS) Nrrrr7rj)Z ignore_indexri)r start_logits end_logitsrt) rSrrr splitr contiguouslenrOclampr rrt)rRrYr9r6rZr r rrrurrr r rZ ignored_indexrZ start_lossZend_lossr~r&r&r'r_sB            z FNetForQuestionAnswering.forwardr)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs<    r) rrrrrrrrrr)Ircr dataclassesr functoolsrtypingrrrr Ztorch.utils.checkpointrZtorch.nnrr r utilsr r ZscipyrZ activationsrZmodeling_outputsrrrrrrrrrZmodeling_utilsrZ pytorch_utilsrrZconfiguration_fnetrZ get_loggerr`loggerr(r*r2Moduler3rfrwr}rrrrrrrrrrrrrrrrrrrr__all__r&r&r&r's~     ,     =&    eY>UI[9 D