a h@sdZddlZddlmZddlmZddlmZmZddl Z ddl Z ddl m Z ddl m Z mZmZdd lmZmZerdd lmZdd lmZdd lmZdd lmZmZmZmZmZmZm Z m!Z!m"Z"ddl#m$Z$ddl%m&Z&ddlm'Z'ddl(m)Z)e'*e+Z,ddZ-ddZ.ddZ/Gddde j0Z1Gddde j0Z2Gddde j0Z3Gdd d e j0Z4Gd!d"d"e j0Z5Gd#d$d$e j0Z6Gd%d&d&eZ7Gd'd(d(e j0Z8Gd)d*d*e j0Z9Gd+d,d,e j0Z:Gd-d.d.e j0Z;Gd/d0d0e j0ZeGd5d6d6e$Z?eed7d8Gd9d:d:eZ@eGd;d<dd?d?e?ZBeGd@dAdAe?ZCedBd8GdCdDdDe?ZDedEd8GdFdGdGe?ZEeGdHdIdIe?ZFeGdJdKdKe?ZGeGdLdMdMe?ZHgdNZIdS)OzPyTorch FNet model.N) dataclass)partial)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)auto_docstringis_scipy_available)linalg)ACT2FN)GradientCheckpointingLayer) 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&b/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/fnet/modeling_fnet.py_two_dim_matmul8s  r(cCs t|||SN)r()r"r#r$r&r&r'two_dim_matmulAsr*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'fftnFs r1cs*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_idxeps 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=Ys zFNetEmbeddings.__init__Nc Cs|dur|}n|dd}|d}|durH|jddd|f}|durt|dr|jddd|f}||d|}|}ntj|tj|jjd}|dur| |}| |} || } | |} | | 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'forwardos,        zFNetEmbeddings.forward)NNNN)__name__ __module__ __qualname____doc__r=r_ __classcell__r&r&rTr'r2Vs r2cs,eZdZfddZddZddZZS)FNetBasicFourierTransformcst||dSr))r<r=_init_fourier_transformrQrTr&r'r=s z"FNetBasicFourierTransform.__init__cCs|jsttjjdd|_n~|jdkrtr|dtj t |j tj d|dtj t |jtj dtt|j|jd|_qtdt|_nt|_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 r/r1fourier_transformrBr rLZtensorr Zdftr@r!tpu_short_seq_lengthr*rlrkrwarningrQr&r&r'rfs$  z1FNetBasicFourierTransform._init_fourier_transformcCs||j}|fSr))rnreal)rR hidden_statesoutputsr&r&r'r_s z!FNetBasicFourierTransform.forward)r`rarbr=rfr_rdr&r&rTr'res recs$eZdZfddZddZZS)FNetBasicOutputcs"ttj|j|jd|_dSNr4)r<r=rrEr@rFrQrTr&r'r=s zFNetBasicOutput.__init__cCs|||}|Sr))rErRrr input_tensorr&r&r'r_szFNetBasicOutput.forwardr`rarbr=r_rdr&r&rTr'rts rtcs$eZdZfddZddZZS)FNetFourierTransformcs"tt||_t||_dSr))r<r=rerRrtoutputrQrTr&r'r=s  zFNetFourierTransform.__init__cCs$||}||d|}|f}|SNr)rRrz)rRrrZ self_outputsfourier_outputrsr&r&r'r_s zFNetFourierTransform.forwardrxr&r&rTr'rys rycs0eZdZfddZejejdddZZS)FNetIntermediatecsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr)) r<r=rrGr@intermediate_sizedense isinstance hidden_actstrrintermediate_act_fnrQrTr&r'r=s   zFNetIntermediate.__init__rrreturncCs||}||}|Sr))rrrRrrr&r&r'r_s  zFNetIntermediate.forwardr`rarbr=r Tensorr_rdr&r&rTr'r}s r}cs4eZdZfddZejejejdddZZS) FNetOutputcsBtt|j|j|_tj|j|jd|_t |j |_ dSru) r<r=rrGr~r@rrErFrIrJrKrQrTr&r'r=s zFNetOutput.__init__)rrrwrcCs&||}||}|||}|Sr))rrKrErvr&r&r'r_s  zFNetOutput.forwardrr&r&rTr'rs rcs,eZdZfddZddZddZZS) FNetLayercs:t|j|_d|_t||_t||_t||_ dSNr) r<r=chunk_size_feed_forward seq_len_dimryfourierr} intermediaterrzrQrTr&r'r=s    zFNetLayer.__init__cCs0||}|d}t|j|j|j|}|f}|Sr{)rrfeed_forward_chunkrr)rRrrZself_fourier_outputsr| layer_outputrsr&r&r'r_s zFNetLayer.forwardcCs||}|||}|Sr))rrz)rRr|Zintermediate_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' z(FNetEncoder.__init__..F) r<r=rSrZ ModuleListr-Znum_hidden_layerslayerZgradient_checkpointingrQrTrr'r= s  zFNetEncoder.__init__FTcCsr|rdnd}t|jD]&\}}|r,||f}||}|d}q|rL||f}|sftdd||fDSt||dS)Nr&rcss|]}|dur|VqdSr)r&)rvr&r&r' rz&FNetEncoder.forward..)last_hidden_staterr) enumeratertupler)rRrroutput_hidden_states return_dictZall_hidden_statesiZ layer_moduleZ layer_outputsr&r&r'r_s    zFNetEncoder.forward)FTrxr&r&rTr'r s rcs0eZdZfddZejejdddZZS) FNetPoolercs*tt|j|j|_t|_dSr))r<r=rrGr@rZTanh activationrQrTr&r'r='s zFNetPooler.__init__rcCs(|dddf}||}||}|Sr{)rr)rRrrZfirst_token_tensor pooled_outputr&r&r'r_,s  zFNetPooler.forwardrr&r&rTr'r&s rcs0eZdZfddZejejdddZZS)FNetPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSru) r<r=rrGr@rrrrrtransform_act_fnrErFrQrTr&r'r=7s   z$FNetPredictionHeadTransform.__init__rcCs"||}||}||}|Sr))rrrErr&r&r'r_@s   z#FNetPredictionHeadTransform.forwardrr&r&rTr'r6s rcs2eZdZfddZddZddddZZS) FNetLMPredictionHeadcsHtt||_t|j|j|_t t |j|_ |j |j_ dSr)) r<r=r transformrrGr@r?decoder Parameterr rNbiasrQrTr&r'r=Hs   zFNetLMPredictionHead.__init__cCs||}||}|Sr))rrrr&r&r'r_Ss  zFNetLMPredictionHead.forwardN)rcCs*|jjjjdkr|j|j_n |jj|_dS)Nmeta)rrrWrrRr&r&r' _tie_weightsXs z!FNetLMPredictionHead._tie_weights)r`rarbr=r_rrdr&r&rTr'rGs rcs$eZdZfddZddZZS)FNetOnlyMLMHeadcstt||_dSr))r<r=r predictionsrQrTr&r'r=bs zFNetOnlyMLMHead.__init__cCs||}|Sr))r)rRsequence_outputprediction_scoresr&r&r'r_fs zFNetOnlyMLMHead.forwardrxr&r&rTr'ras rcs$eZdZfddZddZZS)FNetOnlyNSPHeadcstt|jd|_dSNrg)r<r=rrGr@seq_relationshiprQrTr&r'r=ms zFNetOnlyNSPHead.__init__cCs||}|Sr))r)rRrseq_relationship_scorer&r&r'r_qs zFNetOnlyNSPHead.forwardrxr&r&rTr'rls rcs$eZdZfddZddZZS)FNetPreTrainingHeadscs(tt||_t|jd|_dSr)r<r=rrrrGr@rrQrTr&r'r=xs  zFNetPreTrainingHeads.__init__cCs||}||}||fSr))rr)rRrrrrr&r&r'r_}s  zFNetPreTrainingHeads.forwardrxr&r&rTr'rws rc@s&eZdZUeed<dZdZddZdS)FNetPreTrainedModelrSfnetTcCst|tjr:|jjjd|jjd|jdur|jj nft|tj rz|jjjd|jjd|j dur|jj|j  n&t|tj r|jj |jj ddS)zInitialize the weightsg)meanZstdNg?)rrrGweightdataZnormal_rSZinitializer_rangerZzero_r>r3rEZfill_)rRmoduler&r&r' _init_weightss      z!FNetPreTrainedModel._init_weightsN)r`rarbr__annotations__Zbase_model_prefixZsupports_gradient_checkpointingrr&r&r&r'rs rz0 Output type of [`FNetForPreTraining`]. )Z custom_introc@s^eZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dS)FNetForPreTrainingOutputa 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). Nlossprediction_logitsseq_relationship_logitsrr) r`rarbrcrrr FloatTensorrrrrrrr&r&r&r'rs  rc szeZdZdZd fdd ZddZddZedee j ee j ee j ee j ee ee e eefd d d ZZS) FNetModelz The model can behave as an encoder, following the architecture described in [FNet: Mixing Tokens with Fourier Transforms](https://huggingface.co/papers/2105.03824) by James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, Santiago Ontanon. TcsDt|||_t||_t||_|r2t|nd|_| dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N) r<r=rSr2r^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_embeddingsszFNetModel.set_input_embeddingsN)rYr9r6rZrrrcCs~|dur |n|jj}|dur |n|jj}|durB|durBtdnD|dur\|}|\}} n*|dur~|dd}|\}} ntd|jjr| dkr|jj| krtd|dur|jn|j} |durt|j dr|j j ddd| f} | || } | }nt j |t j| d}|j ||||d} |j| ||d }|d }|jdurP||nd}|sn||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_embedsrjzThe `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_outputrr)rSruse_return_dict ValueErrorrOrmrorWrXr^r9rMr rNrPrrrrr)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'rs( 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. c seZdZddgZfddZddZddZed ee j ee j ee j ee j ee j ee j ee ee e e efd d d ZZS)FNetForPreTrainingcls.predictions.decoder.biascls.predictions.decoder.weightcs,t|t||_t||_|dSr))r<r=rrrclsrrQrTr&r'r="s   zFNetForPreTraining.__init__cCs |jjjSr)rrrrr&r&r'get_output_embeddings+sz(FNetForPreTraining.get_output_embeddingscCs||jj_|j|jj_dSr)rrrrrRZnew_embeddingsr&r&r'set_output_embeddings.s z(FNetForPreTraining.set_output_embeddingsN) rYr9r6rZlabelsnext_sentence_labelrrrc Cs|dur |n|jj}|j||||||d} | dd\} } || | \} } d}|dur|durt}|| d|jj|d}|| dd|d}||}|s| | f| dd}|dur|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 ```Nr9r6rZrrrgr7)rrrrr) rSrrrrviewr?rrr)rRrYr9r6rZrrrrrsrrrr total_lossloss_fctmasked_lm_lossnext_sentence_lossrzr&r&r'r_2s4% zFNetForPreTraining.forward)NNNNNNNN)r`rarb_tied_weights_keysr=rrr rr rrrrrr_rdr&r&rTr'rs0  rc seZdZddgZfddZddZddZed ee j ee j ee j ee j ee j ee ee e e efd d d ZZS)FNetForMaskedLMrrcs,t|t||_t||_|dSr))r<r=rrrrrrQrTr&r'r=|s   zFNetForMaskedLM.__init__cCs |jjjSr)rrr&r&r'rsz%FNetForMaskedLM.get_output_embeddingscCs||jj_|j|jj_dSr)rrr&r&r'rs z%FNetForMaskedLM.set_output_embeddingsNrYr9r6rZrrrrcCs|dur |n|jj}|j||||||d}|d} || } d} |durjt} | | d|jj|d} |s| f|dd} | dur| 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]`. Nrrr7rgrlogitsrr) rSrrrrrr?rrr)rRrYr9r6rZrrrrsrrrrrzr&r&r'r_s&  zFNetForMaskedLM.forward)NNNNNNN)r`rarbrr=rrr rr rrrrrr_rdr&r&rTr'rxs,  rzT FNet Model with a `next sentence prediction (classification)` head on top. c sleZdZfddZedeejeejeejeejeejeeeee e e fdddZ Z S)FNetForNextSentencePredictioncs,t|t||_t||_|dSr))r<r=rrrrrrQrTr&r'r=s   z&FNetForNextSentencePrediction.__init__NrcKsd|vrtdt|d}|dur*|n|jj}|j||||||d} | d} || } d} |durt} | | dd| d} |s| f| dd}| dur| 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.Nrrr7rgr) warningswarn FutureWarningpoprSrrrrrrrr)rRrYr9r6rZrrrkwargsrsrZseq_relationship_scoresrrrzr&r&r'r_s:$   z%FNetForNextSentencePrediction.forward)NNNNNNN)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs&  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. c sleZdZfddZedeejeejeejeejeejeeeee e e fdddZ Z S)FNetForSequenceClassificationcsJt||j|_t||_t|j|_t |j |j|_ | dSr) r<r= num_labelsrrrrIrJrKrGr@ classifierrrQrTr&r'r=s   z&FNetForSequenceClassification.__init__NrcCsr|dur |n|jj}|j||||||d}|d} || } || } d} |dur2|jjdur|jdkrtd|j_n4|jdkr|jtj ks|jtj krd|j_nd|j_|jjdkrt } |jdkr| | | } n | | |} nN|jjdkrt } | | d|j|d} n|jjdkr2t} | | |} |sb| 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_classificationr7rgr)rSrrrKrZ problem_typerr;r rPintr squeezerrrrrr)rRrYr9r6rZrrrrsrrrrrzr&r&r'r_#sF       "     z%FNetForSequenceClassification.forward)NNNNNNN)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs&  rc sleZdZfddZedeejeejeejeejeejeeeee e e fdddZ Z S)FNetForMultipleChoicecs@t|t||_t|j|_t|j d|_ | dSr) r<r=rrrrIrJrKrGr@rrrQrTr&r'r=bs   zFNetForMultipleChoice.__init__NrcCsL|dur |n|jj}|dur&|jdn|jd}|durJ|d|dnd}|durh|d|dnd}|dur|d|dnd}|dur|d|d|dnd}|j||||||d} | d} || } || } | d|} d} |dur t}|| |} |s<| f| dd}| dur8| 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) Nrr7rrgr) rSrrrrOrrKrrrrr)rRrYr9r6rZrrrZ num_choicesrsrrZreshaped_logitsrrrzr&r&r'r_ls:)      zFNetForMultipleChoice.forward)NNNNNNN)r`rarbr=r rr rrrrrr_rdr&r&rTr'r`s&  rc sleZdZfddZedeejeejeejeejeejeeeee e e fdddZ Z S)FNetForTokenClassificationcsJt||j|_t||_t|j|_t |j |j|_ | dSr)rrQrTr&r'r=s   z#FNetForTokenClassification.__init__NrcCs|dur |n|jj}|j||||||d}|d} || } || } d} |durrt} | | d|j|d} |s| f|dd} | dur| 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]`. Nrrr7rgr) rSrrrKrrrrrrr)rRrYr9r6rZrrrrsrrrrrzr&r&r'r_s(   z"FNetForTokenClassification.forward)NNNNNNN)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs&  rc steZdZfddZedeejeejeejeejeejeejeeeee e e fd ddZ Z S)FNetForQuestionAnsweringcs<t||j|_t||_t|j|j|_| dSr)) r<r=rrrrrGr@ qa_outputsrrQrTr&r'r=s   z!FNetForQuestionAnswering.__init__N) rYr9r6rZstart_positions end_positionsrrrc CsB|dur |n|jj}|j||||||d} | d} || } | jddd\} } | d} | d} d}|dur|durt|dkr|d}t|dkr|d}| d}| d|}| d|}t |d}|| |}|| |}||d}|s0| | f| dd}|dur,|f|S|St || | | j dS) Nrrrr7rh)Z ignore_indexrg)r start_logits end_logitsrr) rSrrrsplitr contiguouslenrOclamprrrr)rRrYr9r6rZrrrrrsrrrrrZ ignored_indexrZ start_lossZend_lossrzr&r&r'r_sB             z FNetForQuestionAnswering.forward)NNNNNNNN)r`rarbr=r rr rrrrrr_rdr&r&rTr'rs*  r) rrrrrrrrrr)Jrcr dataclassesr functoolsrtypingrrr Ztorch.utils.checkpointrZtorch.nnrrr utilsr r Zscipyr Z activationsrZmodeling_layersrZmodeling_outputsrrrrrrrrrZmodeling_utilsrZ pytorch_utilsrrZconfiguration_fnetrZ get_loggerr`loggerr(r*r1Moduler2rertryr}rrrrrrrrrrrrrrrrrrr__all__r&r&r&r's      ,      =&    eY>UI[9D