rhddlZddlZddlZddlmZddlmZmZddlZddlm Z ddl m Z m Z m Z ddlmZddlmZmZmZmZmZmZmZmZdd lmZdd lmZmZdd lmZm Z m!Z!d d l"m#Z#e!jHe%Z&dZ'Gdde jPZ)e jTe)dZ+Gdde jPZ,Gdde jPZ-Gdde jPZ.Gdde jPZ/Gdde jPZ0Gdde jPZ1Gdde jPZ2Gd d!e jPZ3Gd"d#e jPZ4Gd$d%e jPZ5Gd&d'e jPZ6Gd(d)e jPZ7Gd*d+e jPZ8Gd,d-e jPZ9Gd.d/e jPZ:Gd0d1e jPZ;Gd2d3e jPZ<Gd4d5e jPZ=e Gd6d7eZ>ee d89Gd:d;eZ?e Gd<d=e>Z@e d>9Gd?d@e>ZAe GdAdBe>ZBGdCdDe jPZCe dE9GdFdGe>ZDe dH9GdIdJe>ZEe GdKdLe>ZFe GdMdNe>ZGe GdOdPe>ZHgdQZIy)RN) dataclass)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BaseModelOutputBaseModelOutputWithPoolingMaskedLMOutputMultipleChoiceModelOutputNextSentencePredictorOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputauto_docstringlogging)MobileBertConfigc  ddl}ddl}ddl}t jj|}tjd||jj|}g}g} |D]^\} } tjd| d| |jj|| } |j| | j| `t|| D]\} } | jdd} | jd d } | jd d } | jd d} | j!d} t#d| Dr(tjddj%| |} | D]}|j'd|r|j!d|}n|g}|ddk(s|ddk(r t)| d} nW|ddk(s|ddk(r t)| d} n:|ddk(r t)| d} n%|ddk(r t)| d} n t)| |d} t-|dk\st/|d}| |} ddd k(r t)| d} n|dk(r|j1| } | j2| j2k(s"Jd!| j2d"| j2d# tjd$| t9j:| | _|S#t$rtj dwxYw#t*$r+tjddj%| YwxYw#t4$r1}|xj6| j2| j2fz c_d}~wwxYw)%z'Load tf checkpoints in a pytorch model.rNzLoading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.z&Converting TensorFlow checkpoint from zLoading TF weight z with shape ffn_layerffn FakeLayerNorm LayerNormextra_output_weightsz dense/kernelbert mobilebert/c3$K|]}|dv yw))adam_vadam_mAdamWeightDecayOptimizerAdamWeightDecayOptimizer_1 global_stepN).0ns /var/www/html/ai-insurance-compliance-backend/venv/lib/python3.12/site-packages/transformers/models/mobilebert/modeling_mobilebert.py z0load_tf_weights_in_mobilebert..Vs  n n sz Skipping z [A-Za-z]+_\d+z_(\d+)kernelgammaweight output_biasbetabiasoutput_weightssquad classifierri _embeddingszPointer shape z and array shape z mismatchedzInitialize PyTorch weight )renumpy tensorflow ImportErrorloggererrorospathabspathinfotrainlist_variables load_variableappendzipreplacesplitanyjoin fullmatchgetattrAttributeErrorlenint transposeshapeAssertionErrorargstorch from_numpydata)modelconfigtf_checkpoint_pathr;nptftf_path init_varsnamesarraysnamerTarraypointerm_name scope_namesnumes r.load_tf_weights_in_mobilebertrj5sh ggoo01G KK8 BC''0I E F  e (l5'BC&&w5 T e  5&)1/ e||K/||O[9||2NC||FL1zz#     KK)CHHTN#34 5  'F||,f5 hhy&9 %h 1~)[^w-F!'84Q=0KNf4L!'62Q#33!'84Q7*!'<8%g{1~>G;1$+a.)!#,+ ', #$<= (gx0G x LL'E ==EKK/  /@ [Y /  078''. c1/d LI   Q   b&KK)CHHTN+; <=  FFw}}ekk2 2F  s5 KK6 ;L- K360L*)L*- M'6,M""M'cXeZdZdfd Zdej dej fdZxZS)NoNormct|tjt j ||_tjt j||_yN) super__init__r ParameterrWzerosr5onesr2)self feat_sizeeps __class__s r.rpzNoNorm.__init__s@ LLY!78 ll5::i#89 input_tensorreturnc:||jz|jzSrn)r2r5)rtrys r.forwardzNoNorm.forwardsdkk)DII55rxrn__name__ __module__ __qualname__rprWTensorr| __classcell__rws@r.rlrls#: 6ELL6U\\6rxrl) layer_normno_normc eZdZdZfdZ d deejdeejdeejdeejdejf dZ xZ S) MobileBertEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.cXt||j|_|j|_|j|_t j |j|j|j|_ t j |j|j|_ t j |j|j|_ |jrdnd}|j|z}t j||j|_t!|j"|j|_t j&|j(|_|j-dt/j0|jj3ddy)N) padding_idxr r position_ids)rF) persistent)rorp trigram_inputembedding_size hidden_sizer Embedding vocab_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddingsLinearembedding_transformationNORM2FNnormalization_typer Dropouthidden_dropout_probdropoutregister_bufferrWarangeexpand)rtr[embed_dim_multiplierembedded_input_sizerws r.rpzMobileBertEmbeddings.__init__sF #11$33!--!||F,=,=v?T?Tbhbubuv#%<<0N0NPVPbPb#c %'\\&2H2H&J\J\%]"$($6$6qA"114HH(* 2EvGYGY(Z% !:!:;F J"%88;PP ^^J/ \\*- rx)NNNN) r~rr__doc__rprrW LongTensor FloatTensorrr|rrs@r.rrs~Q 015593759 0E,,-0!!1!120u//0 0   1 12 0  0rxrceZdZfdZ d dej dej dej deejdeejdeede ej fd Z xZ S) MobileBertSelfAttentionc`t||j|_t|j|jz |_|j|j z|_tj|j|j |_ tj|j|j |_ tj|jr |jn |j|j |_ tj|j|_yrn)rorpnum_attention_headsrRtrue_hidden_sizeattention_head_size all_head_sizerrquerykeyuse_bottleneck_attentionrrrattention_probs_dropout_probrrtr[rws r.rpz MobileBertSelfAttention.__init__s #)#=#= #&v'>'>A[A['[#\ !558P8PPYYv668J8JK 99V44d6H6HIYY'-'F'FF # #FL^L^`d`r`r  zz&"E"EF rx query_tensor key_tensor value_tensorattention_mask head_maskoutput_attentionsrzc|j\}}} |j|j|d|j|jj dd} |j |j|d|j|jj dd} |j|j|d|j|jj dd} tj| | j dd} | tj|jz } || |z} tjj| d}|j|}|||z}tj|| }|j!ddddj#}|j%dd|j&fz}|j|}|r||f}|S|f}|S)Nrrr9rrr )rTrviewrrrSrrrWmatmulmathsqrtrrsoftmaxrpermute contiguousrr)rtrrrrrr batch_sizer_ query_layer key_layer value_layerattention_scoresattention_probs context_layernew_context_layer_shapeoutputss r.r|zMobileBertSelfAttention.forwards%1$6$6! J JJ| $ T*b$":":DS"S%**+BC 6G=/2O\M]rxNNN r~rrrprWrrrbooltupler|rrs@r.rrs G$7;15,0-ll-LL-ll - !!2!23 - E--. -$D>- u|| -rxrcneZdZfdZdej dej dej fdZxZS)MobileBertSelfOutputcjt||j|_tj|j |j |_t|j|j |j|_ |js%tj|j|_ yyNrv)rorpuse_bottleneckrrrdenserrlayer_norm_epsr rrrrs r.rpzMobileBertSelfOutput.__init__s $33YYv668O8OP  !:!:;Ft|dk(ryt||jj|jj|j \}}t |jj||j_t |jj||j_t |jj||j_ t |jj|d|j_ |jjt|z |j_|jj|jjz|j_ |j j||_y)Nrrr)rQrrtrrrrrrrrrrunion)rtheadsindexs r. prune_headszMobileBertAttention.prune_heads2s u:? 7 49900$))2O2OQUQbQb u -TYY__eD *499==%@ ,TYY__eD .t{{/@/@%QO )- (E(EE (R %"&))"?"?$))B_B_"_  --33E:rxrrr layer_inputrrrrzcn|j||||||}|j|d|} | f|ddz} | S)Nrr)rtr) rtrrrrrrr self_outputsattention_outputrs r.r|zMobileBertAttention.forwardDsTyy         ;;|A D#% QR(88rxr) r~rrrprrWrrrrrr|rrs@r.rr+s" ;07;15,0llLLll  \\  !!2!23 E--.$D> u|| rxrcVeZdZfdZdej dej fdZxZS)MobileBertIntermediatect|tj|j|j |_t|jtrt|j|_ y|j|_ yrn) rorprrrintermediate_sizer isinstance hidden_actstrr intermediate_act_fnrs r.rpzMobileBertIntermediate.__init__^s] YYv668P8PQ f'' -'-f.?.?'@D $'-'8'8D $rxrrzcJ|j|}|j|}|Srn)rrrtrs r.r|zMobileBertIntermediate.forwardfs& =1 00? rxr}rs@r.rr]s#9U\\ellrxrcneZdZfdZdej dej dej fdZxZS)OutputBottleneckc.t|tj|j|j |_t|j|j |j|_ tj|j|_ yr)rorprrrrrrrrr rrrrs r.rpzOutputBottleneck.__init__msh YYv668J8JK  !:!:;F U\\ELLUZUaUarxr ceZdZfdZdej dej dej dej fdZxZS)MobileBertOutputcrt||j|_tj|j |j |_t|j|j |_ |js%tj|j|_ yt||_yrn)rorprrrrrrrrr rrrr  bottleneckrs r.rpzMobileBertOutput.__init__{s $33YYv779P9PQ  !:!:;F>,9J*JKL >>,9J*JKL??<9JKLrxr}rs@r.rrzs?7 #(<< DILL ejeqeq   rxrcVeZdZfdZdej dej fdZxZS)BottleneckLayerct|tj|j|j |_t|j|j |j|_ yr) rorprrrintra_bottleneck_sizerrrrr rs r.rpzBottleneckLayer.__init__sR YYv1163O3OP  !:!:;F"  ( (.014 4  - -%)^^M%B "*,BMSmn n!=-A[\ \rx r~rrrprWrrr|rrs@r.rrs+5]U\\]eELL6I]rxrcneZdZfdZdej dej dej fdZxZS) FFNOutputct|tj|j|j |_t|j|j |j|_ yr) rorprrrrrrrrr rs r.rpzFFNOutput.__init__sR YYv779P9PQ  !:!:;F $7G rxr}rs@r.r,r,s#( U\\ellrxr,c eZdZfdZ ddej deejdeejdeede ej f dZ xZ S) MobileBertLayerct||j|_|j|_t ||_t ||_t||_ |jrt||_ |jdkDrHtjt|jdz Dcgc] }t|c}|_yycc}wNr)rorprnum_feedforward_networksrr!rr.rrrrr ModuleListranger,rrtr[rrws r.rpzMobileBertLayer.__init__s $33(.(G(G%,V426:&v.   (0DO  * *Q .}}fFeFehiFi@j%k1hv&6%klDH /%ks6Crrrrrzc |jr|j|\}}}}n |gdz\}}}}|j|||||||} | d} | f} | dd} |jdk7r+t |j D]\} }|| } | | fz } |j | }|j|| |}|f| ztjd||||| |fz| z} | S)Nr#)rrri) rrr!r5 enumeraterr.rrWtensor)rtrrrrrrrrself_attention_outputsrsri ffn_moduler0rs r.r|zMobileBertLayer.forwards>   BF//R_B` ?L*lKCP/TUBU ?L*lK!%      /"0" 2!4  (,  ( (A -!*488!4 ) :#-.>#? &(( )#//0@A{{#68H-X O  T" #   rxrrrs@r.r2r2sp m 7;15,0 .||.!!2!23.E--. . $D> . u||  .rxr2ceZdZfdZ d dej deejdeejdeedeedeede e e ffd Z xZ S) MobileBertEncoderct|tjt |j Dcgc] }t |c}|_ycc}wrn)rorprr6r7num_hidden_layersr2layerr8s r.rpzMobileBertEncoder.__init__s< ]]U6KcKcEd#eOF$;#ef #esArrrroutput_hidden_states return_dictrzc|rdnd}|rdnd}t|jD],\} } |r||fz}| |||| |} | d}|s$|| dfz}.|r||fz}|std|||fDSt|||S)Nr+rrc3&K|] }|| ywrnr+)r,vs r.r/z,MobileBertEncoder.forward..=shqZ[Zghs)last_hidden_stater attentions)r:rDrr ) rtrrrrrErFall_hidden_statesall_attentionsr> layer_modulers r.r|zMobileBertEncoder.forwards#7BD0d(4 FOA|#$58H$H!(! ! M *!,M !/=3C2E!E F 1]4D D h]4E~$Vhh h+;LYg  rx)NNFFT)r~rrrprWrrrrrrr r|rrs@r.rArAsg7;15,1/4&*" ||" !!2!23" E--. " $D> " 'tn " d^"  uo% &" rxrAcVeZdZfdZdej dej fdZxZS)MobileBertPoolerct||j|_|jr0t j |j |j |_yyrn)rorpclassifier_activation do_activaterrrrrs r.rpzMobileBertPooler.__init__DsH !77   6#5#5v7I7IJDJ rxrrzc|dddf}|js|S|j|}tj|}|S)Nr)rSrrWtanh)rtrfirst_token_tensor pooled_outputs r.r|zMobileBertPooler.forwardJsE+1a40% % JJ'9:M!JJ}5M rxr}rs@r.rPrPCs$K !U\\ !ell !rxrPcVeZdZfdZdej dej fdZxZS)!MobileBertPredictionHeadTransformcZt|tj|j|j|_t |jtrt|j|_ n|j|_ td|j|j|_ y)Nrr)rorprrrrrrrr transform_act_fnrrr rs r.rpz*MobileBertPredictionHeadTransform.__init__Ws| YYv1163E3EF f'' -$*6+<+<$=D !$*$5$5D ! .v/A/AvG\G\]rxrrzcl|j|}|j|}|j|}|Srn)rr[r rs r.r|z)MobileBertPredictionHeadTransform.forward`s4 =1 --m< }5 rxr}rs@r.rYrYVs$^U\\ellrxrYc^eZdZfdZddZdej dej fdZxZS)MobileBertLMPredictionHeadct|t||_t j |j |j|jz d|_ t j |j|j d|_ t jtj|j |_|j|j_y)NF)r5)rorprY transformrrrrrrdecoderrqrWrrr5rs r.rpz#MobileBertLMPredictionHead.__init__hs :6BYYv00&2D2DvG\G\2\chi yy!6!68I8IPUV LLV->->!?@  II rxrzc:|j|j_yrn)r5rarts r. _tie_weightsz'MobileBertLMPredictionHead._tie_weightsss II rxrc|j|}|jtj|jj j |jj gd}||jjz }|S)Nrr) r`rrWrrar2trr5rs r.r|z"MobileBertLMPredictionHead.forwardvsk}5 %,,UYY 8K8K8M8M8OQUQ[Q[QbQb7cij-kl *** rx)rzN) r~rrrprdrWrr|rrs@r.r^r^gs( &&U\\ellrxr^cVeZdZfdZdej dej fdZxZS)MobileBertOnlyMLMHeadcBt|t||_yrn)rorpr^ predictionsrs r.rpzMobileBertOnlyMLMHead.__init__~s 5f=rxsequence_outputrzc(|j|}|Srn)rj)rtrkprediction_scoress r.r|zMobileBertOnlyMLMHead.forwards ,,_=  rxr}rs@r.rhrh}s#>!u||! !rxrhcteZdZfdZdej dej deej fdZxZS)MobileBertPreTrainingHeadsct|t||_t j |j d|_yNr9)rorpr^rjrrrseq_relationshiprs r.rpz#MobileBertPreTrainingHeads.__init__s4 5f= " &*<*. v 6 7 KK   " " $ MM   $ $S )  : ; KK   " " $cd d  "  6 6y. Q#..*K  &',,.s3KTU U%.%:!! @T@T  !"ZZ FCN  !"[[EJJvVN150P0PQ_al0m&&y$++2O2OP ??l>iv+ ,, 2/!5# ' *!,8< 8O O4UY #]3oab6II I)-')77&11   rx)T) NNNNNNNNN)r~rrrrprrrrrrWrrrrrr r|rrs@r.rrs /0C156:59371559/3,0&*D E,,-D !!2!23D !!1!12 D u//0 D E--. D   1 12D 'tnD $D>D d^D  u00 1D D rxrz MobileBert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next sentence prediction (classification)` head. ceZdZddgZfdZdZdZddeede jffd Z e dd ee jd ee jd ee jd ee jd ee jdee jdee jdee jdee jdee jdee jdeeeffdZxZS)MobileBertForPreTrainingcls.predictions.decoder.weightcls.predictions.decoder.biasct||t||_t ||_|j yrn)rorprr#roclsrrs r.rpz!MobileBertForPreTraining.__init__9s4  )&1-f5 rxcB|jjjSrnrrjrarcs r.get_output_embeddingsz.MobileBertForPreTraining.get_output_embeddingsAxx##+++rxc||jj_|j|jj_yrnrrjrar5rtnew_embeddingss r.set_output_embeddingsz.MobileBertForPreTraining.set_output_embeddingsD,'5$$2$7$7!rxnew_num_tokensrzc|j|jjj|d|jj_t||SNT)r transposed)r_get_resized_lm_headrrjrroresize_token_embeddingsrtrrws r.rz0MobileBertForPreTraining.resize_token_embeddingsHsR%)%>%> HH & &~RV&?& "w.n.MMrxrrrrrrlabelsnext_sentence_labelrrErFc  | | n|jj} |j||||||| | |  } | dd\} }|j| |\}}d}|u|st }||j d|jj |j d}||j dd|j d}||z}| s||f| ddz}||f|zS|St|||| j| jS)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]` 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. Examples: ```python >>> from transformers import AutoTokenizer, MobileBertForPreTraining >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased") >>> model = MobileBertForPreTraining.from_pretrained("google/mobilebert-uncased") >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) >>> # Batch size 1 >>> outputs = model(input_ids) >>> prediction_logits = outputs.prediction_logits >>> seq_relationship_logits = outputs.seq_relationship_logits ```NrrrrrrrErFr9r)rrrrrK) r[rr#rrrrrrrK)rtrrrrrrrrrrErFrrkrWrmrt total_lossloss_fctmasked_lm_lossnext_sentence_lossrs r.r|z MobileBertForPreTraining.forwardPsGV&1%9=<@37K E,,-K !!2!23K !!1!12 K u//0 K E--. K   1 12K ))*K &e&6&67K $E$5$56K 'u'8'89K e//0K  u44 5K K rxrceZdZddgZfdZdZdZddeede jffd Z e dd ee jd ee jd ee jd ee jd ee jdee jdee jdeedeedeedeeeffdZxZS)MobileBertForMaskedLMrrct||t|d|_t ||_||_|jyNF)r)rorprr#rhrr[rrs r.rpzMobileBertForMaskedLM.__init__s=  )&EJ(0  rxcB|jjjSrnrrcs r.rz+MobileBertForMaskedLM.get_output_embeddingsrrxc||jj_|j|jj_yrnrrs r.rz+MobileBertForMaskedLM.set_output_embeddingsrrxrrzc|j|jjj|d|jj_t||Srrrs r.rz-MobileBertForMaskedLM.resize_token_embeddingssR%)%>%> HH & &~RV&?& "w.n.MMrxrrrrrrrrrErFc | | n|jj} |j|||||||| |  } | d} |j| } d}|Ft }|| j d|jj |j d}| s| f| ddz}||f|zS|St|| | j| jS)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]` Nrrrr9rlogitsrrK) r[rr#rrrrrrrK)rtrrrrrrrrrErFrrkrmrrrs r.r|zMobileBertForMaskedLM.forwards(&1%2 'tn2 d^2  un$ %2 2 rxrcVeZdZfdZdej dej fdZxZS)MobileBertOnlyNSPHeadclt|tj|jd|_yrq)rorprrrrrrs r.rpzMobileBertOnlyNSPHead.__init__s'  " &*<*>> from transformers import AutoTokenizer, MobileBertForNextSentencePrediction >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased") >>> model = MobileBertForNextSentencePrediction.from_pretrained("google/mobilebert-uncased") >>> 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])) >>> loss = outputs.loss >>> logits = outputs.logits ```rzoThe `next_sentence_label` argument is deprecated and will be removed in a future version, use `labels` instead.Nrrrr9r) warningswarn FutureWarningpopr[rr#rrrrrrK)rtrrrrrrrrrErFkwargsrrWrtrrrs r.r|z+MobileBertForNextSentencePrediction.forward sR !F * MM%  ZZ 56F%0%O 'tnO d^O  u11 2O O rxrz MobileBert 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^eZdZfdZe ddeejdeejdeejdeejdeejdeejdeejd eed eed eed e e eje ffd Z xZ S)#MobileBertForSequenceClassificationcnt|||j|_||_t ||_|j |j n |j}tj||_ tj|j|j|_ |jyrn)rorp num_labelsr[rr#classifier_dropoutrrrrrrr8rrtr[rrws r.rpz,MobileBertForSequenceClassification.__init__ds   ++ )&1)/)B)B)NF % %TZTnTn zz"45 ))F$6$68I8IJ rxrrrrrrrrrErFrzc @| | n|jj} |j|||||||| |  } | d} |j| } |j | } d}||jj |j dk(rd|j_nl|j dkDrL|jtjk(s|jtjk(rd|j_nd|j_|jj dk(rIt}|j dk(r&|| j|j}n|| |}n|jj dk(r=t}|| jd|j |jd}n,|jj dk(rt}|| |}| s| f| ddz}||f|zS|St!|| | j"| j$ 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). Nrr regressionsingle_label_classificationmulti_label_classificationrr9r)r[rr#rr8 problem_typerrrWrrRr squeezerrrrrrK)rtrrrrrrrrrErFrrWrrrrs r.r|z+MobileBertForSequenceClassification.forwardss(&1%E 'tnE d^E  uU\\"$<< =E E rxrc~eZdZfdZe ddeejdeejdeejdeejdeejdeejdeejd eejd eed eed eed e e eje ffdZ xZ S)MobileBertForQuestionAnsweringct|||j|_t|d|_t j |j|j|_|jyr) rorprrr#rrr qa_outputsrrs r.rpz'MobileBertForQuestionAnswering.__init__sU   ++)&EJ))F$6$68I8IJ rxrrrrrrstart_positions end_positionsrrErFrzc (| | n|jj} |j||||||| | |  } | d} |j| }|j dd\}}|j dj }|j dj }d}||t|jdkDr|j d}t|jdkDr|j d}|jd}|jd|}|jd|}t|}|||}|||}||zdz }| s||f| ddz}||f|zS|St|||| j| jS) Nrrrrr) ignore_indexr9)r start_logits end_logitsrrK)r[rr#rrKrrrQrclamprrrrK)rtrrrrrrrrrrErFrrkrrrr ignored_indexr start_lossend_lossrs r.r|z&MobileBertForQuestionAnswering.forwards&1% ELL)> !.> !. > u||, > ELL) >   -> "%,,/>   -> $D>> 'tn> d^>  uU\\"$@@ A> > rxrc^eZdZfdZe ddeejdeejdeejdeejdeejdeejdeejd eed eed eed e e eje ffd Z xZ S)MobileBertForMultipleChoicec*t||t||_|j |jn |j }t j||_t j|jd|_ |jyr4) rorprr#rrrrrrrr8rrs r.rpz$MobileBertForMultipleChoice.__init__su  )&1)/)B)B)NF % %TZTnTn zz"45 ))F$6$6: rxrrrrrrrrrErFrzc L| | n|jj} ||jdn|jd} |!|jd|j dnd}|!|jd|j dnd}|!|jd|j dnd}|!|jd|j dnd}|1|jd|j d|j dnd}|j |||||||| |  } | d} |j | } |j| }|jd| }d}|t}|||}| s|f| ddz}||f|zS|St||| j| jS)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) Nrrrrr9r) r[rrTrrr#rr8rrrrK)rtrrrrrrrrrErF num_choicesrrWrreshaped_logitsrrrs r.r|z#MobileBertForMultipleChoice.forwardsX&1%G>SINN2y~~b'9:Y] M[Mg,,R1D1DR1HImqM[Mg,,R1D1DR1HImqGSG_|((\->->r-BCei (   r=#5#5b#9=;M;Mb;Q R  // ))%'/!5#"     ]3 / ++b+6  ')HOV4D%''!"+5F)-)9TGf$ Ev E("!//))   rxr)r~rrrprrrWrrrrrr|rrs@r.rr s -11515/3,004)-,0/3&*X ELL)X !.X !. X u||, X ELL) X   -X &X $D>X 'tnX d^X  uU\\"$== >X X rxrc^eZdZfdZe ddeejdeejdeejdeejdeejdeejdeejd eed eed eed e e eje ffd Z xZ S) MobileBertForTokenClassificationcdt|||j|_t|d|_|j |j n |j }tj||_ tj|j|j|_ |jyr)rorprrr#rrrrrrrr8rrs r.rpz)MobileBertForTokenClassification.__init__zs   ++)&EJ)/)B)B)NF % %TZTnTn zz"45 ))F$6$68I8IJ rxrrrrrrrrrErFrzc | | n|jj} |j|||||||| |  } | d} |j| } |j | } d}|2 'tn2 d^2  uU\\"$99 :2 2 rxr) rrrrrrrr2rrvrj)JrrAr dataclassesrtypingrrrWrtorch.nnrrr activationsr modeling_outputsr r rrrrrrmodeling_utilsr pytorch_utilsrrutilsrrrconfiguration_mobilebertr get_loggerr~r?rjModulerlr rrrrrrr rrrr(r,r2rArPrYr^rhrorvrrrrrrrrrr__all__r+rxr.r s'. !" AA!   .Q996   H %K\6RYY6& 9I299IX;bii;|299"/"))/d RYY  ryy ryy0 bii !]!]H   ryy <bii<~'  ' T!ryy!& ",!BII! 9 9%%%0  :[:  :&g /g g T f 8f  f RM 5M M `&BII& Z *CZ  Z zV *CV V rJ %>J J Zg ";g g TB '@B B J rx