o Zhj@sZdZddlZddlZddlmZddlmZmZmZm Z m Z ddl Z ddl m Z ddl mZmZmZddlmZdd lmZdd lmZmZmZmZmZmZmZdd lmZdd lm Z m!Z!m"Z"m#Z#dd l$m%Z%m&Z&m'Z'ddl(m)Z)e'*e+Z,ddZ-Gddde j.Z/Gddde j.Z0Gddde0Z1e0e1dZ2Gddde j.Z3Gddde j.Z4Gddde j.Z5e&Gdd d eZ6eGd!d"d"e%Z7e&Gd#d$d$e6Z8e&d%d&Gd'd(d(e6Z9Gd)d*d*e j.Z:Gd+d,d,e j.Z;e&Gd-d.d.e6Ze&Gd4d5d5e6Z?e&Gd6d7d7e6Z@gd8ZAdS)9zPyTorch ALBERT model.N) dataclass)DictListOptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)#_prepare_4d_attention_mask_for_sdpa)BaseModelOutputBaseModelOutputWithPoolingMaskedLMOutputMultipleChoiceModelOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indices"is_torch_greater_or_equal_than_2_2prune_linear_layer) ModelOutputauto_docstringlogging) AlbertConfigc Cs*zddl}ddl}ddl}Wn tytdwtj|}t d||j |}g}g} |D] \} } t d| d| |j || } | | | | q6t|| D]\} } t| q\t|| D]\} } | } | dd} | d d } | d d } | d d} | dd} | dd} | dd} | dd} | dd} | dd} | dd} | dd} | dd} | dd} | dd} | d d!} | d"d#} | d$d%} t| dd&krd'| vsd(| vrd)| } d*| vr| d+d,} | d-d.} | d} d/| vs!d0| vs!d1| vs!d2| vs!d3| vr-t d4d| qj|}| D]}|d5|rA|d6|}n|g}|dd7ksR|dd8krXt|d.}nN|dd'ksf|dd9krlt|d:}n:|dd(kryt|d.}n-|dd;krt|d<}n z t||d}Wntyt d4d| Yq1wt|d=krt|d&}||}q1|d>dd?krt|d.}n |d7kr|| } z|j| jkrtd@|jdA| jdBWnty}z |j|j| jf7_d}~wwtdC| dD| t| |_qj|S)Ez'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 zmodule/Zffn_1ffnzbert/zalbert/Z attention_1 attentionz transform/Z LayerNorm_1full_layer_layer_norm LayerNormzattention/LayerNormz transformer/zintermediate/dense/zffn/intermediate/output/dense/z ffn_output/z/output//z/self/z pooler/densepoolerzcls/predictions predictionszpredictions/attentionzembeddings/attention embeddingsZ inner_group_zalbert_layers/Zgroup_zalbert_layer_groups/rZ output_biasZoutput_weightsz classifier/Zseq_relationshipzseq_relationship/output_zsop_classifier/classifier/weightsweightZadam_mZadam_vZAdamWeightDecayOptimizerZAdamWeightDecayOptimizer_1Z global_stepz Skipping z [A-Za-z]+_\d+z_(\d+)ZkernelgammabetabiasZsquad classifieriZ _embeddingszPointer shape z and array shape z mismatchedzInitialize PyTorch weight z from )renumpyZ tensorflow ImportErrorloggererrorospathabspathinfotrainZlist_variablesZ load_variableappendzipprintreplacelensplitjoin fullmatchgetattrAttributeErrorint transposeshape ValueErrorargstorchZ from_numpydata)modelconfigZtf_checkpoint_pathr0nptfZtf_pathZ init_varsnamesZarraysnamerFarray original_nameZpointerZm_nameZ scope_namesnumerUY/var/www/auris/lib/python3.10/site-packages/transformers/models/albert/modeling_albert.pyload_tf_weights_in_albert3s                          "                 rWc sneZdZdZdeffdd Z     ddeejdeejd eejd eej d e d ej f d dZ Z S)AlbertEmbeddingszQ Construct the embeddings from word, position and token_type embeddings. rLcsttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_|jdt|jdddt|dd|_|jd tj|jtjd dddS) N) padding_idxZeps position_ids)rF) persistentposition_embedding_typeabsolutetoken_type_idsdtype)super__init__r Embedding vocab_sizeembedding_sizeZ pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingsZtype_vocab_sizetoken_type_embeddingsr$layer_norm_epsDropouthidden_dropout_probdropoutZregister_bufferrIarangeexpandrBr^zerosr[sizelongselfrL __class__rUrVrds  zAlbertEmbeddings.__init__Nr input_idsr`r[ inputs_embedspast_key_values_lengthreturnc Cs|dur |}n|dd}|d}|dur&|jdd|||f}|durPt|drE|jddd|f}||d|} | }n tj|tj|jjd}|durY| |}| |} || } |j dkrp| |} | | 7} | | } || } | S)Nr\rr`rrbdevicer_)rsr[hasattrr`rqrIrrrtr~rhrkr^rjr$ro) rvryr`r[rzr{ input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedrkr(rjrUrUrVforwards,        zAlbertEmbeddings.forward)NNNNr)__name__ __module__ __qualname____doc__rrdrrI LongTensor FloatTensorrDTensorr __classcell__rUrUrwrVrXs*rXcseZdZdeffdd ZdejdejfddZdee dd fd d Z dd ejde ej de ej de deeejeejejfff ddZZS)AlbertAttentionrLcs4t|j|jdkrt|dstd|jd|j|j|_|j|_|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t |j|_t |j|j|_tj|j|jd|_t|_t|dd|_|jdks|jd kr|j|_td |jd |j|_dSdS) NrrgzThe hidden size (z6) is not a multiple of the number of attention heads (rZr^r_ relative_keyrelative_key_queryr/r)rcrd hidden_sizenum_attention_headsrrGattention_head_size all_head_sizerLinearquerykeyvaluermattention_probs_dropout_probattention_dropoutrnoutput_dropoutdenser$rlset pruned_headsrBr^riredistance_embeddingrurwrUrVrds0  zAlbertAttention.__init__xr|cCs6|dd|j|jf}||}|ddddS)Nr\rr/rr )rsrrviewZpermute)rvrZ new_x_shaperUrUrVtranspose_for_scoress z$AlbertAttention.transpose_for_scoresheadsNcCst|dkrdSt||j|j|j\}}t|j||_t|j||_t|j||_t|j |dd|_ |jt||_|j|j|_ |j ||_dS)Nrrdim) r>rrrrrrrrrrunion)rvrindexrUrUrV prune_headss zAlbertAttention.prune_headsF hidden_statesattention_mask head_maskoutput_attentionscCs||}||}||}||}||} ||} t|| dd} | t|j } |dur8| |} |j dksB|j dkr| d} tj | tj |jddd} tj | tj |jddd}| |}|||jd}|j|jd}|j dkrtd||}| |} n|j dkrtd||}td | |}| ||} tjj| dd }||}|dur||}t|| }|d dd }||}||}|||}|r||fS|fS) Nr\rrrr}razbhld,lrd->bhlrzbhrd,lrd->bhlrrr/)rrrrrImatmulrEmathsqrtrr^rsrprtr~rrritorbZeinsumr functionalZsoftmaxrflattenrrr$)rvrrrrZmixed_query_layerZmixed_key_layerZmixed_value_layer query_layer key_layer value_layerZattention_scoresrZposition_ids_lZposition_ids_rZdistanceZpositional_embeddingZrelative_position_scoresZrelative_position_scores_queryZrelative_position_scores_keyZattention_probsZ context_layerprojected_context_layerprojected_context_layer_dropoutlayernormed_context_layerrUrUrVr*sD               zAlbertAttention.forwardNNF)rrrrrdrIrrrrDrrrboolrrrrrUrUrwrVrs$rcsneZdZfddZ   d dejdeejdeejded e e eje ejejfff fd d Z Z S) AlbertSdpaAttentioncs t||j|_t |_dSN)rcrdr dropout_probrrequire_contiguous_qkvrurwrUrVrdfs  zAlbertSdpaAttention.__init__NFrrrrr|cs|jdks|rtdtj|||dS|\}}}|||}|||} || |} |j rM|j j dkrM|durM| }| } | } tjjj|| | ||jr[|jnddd} | dd } | |||j} || } || } ||| }|fS) Nr_aAlbertSdpaAttention is used but `torch.nn.functional.scaled_dot_product_attention` does not support non-absolute `position_embedding_type` or `output_attentions=True` . Falling back to the eager attention implementation, but specifying the eager implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.)rcudaF)rrrZ attn_maskZ dropout_pZ is_causalrr/)r^r3warningrcrrsrrrrrr~type contiguousrIrrZscaled_dot_product_attentionZtrainingrrEZreshaperrrr$)rvrrrr batch_sizeZseq_len_rrrattention_outputrrrrwrUrVrks6   zAlbertSdpaAttention.forwardr) rrrrdrIrrrrrrrrrUrUrwrVres  r)eagersdpacs~eZdZdeffdd Z    ddejdeejdeejd e d e d e ejejff d d Z dejd ejfddZ Z S) AlbertLayerrLcst||_|j|_d|_tj|j|jd|_ t |j ||_ t |j|j|_t |j|j|_t|j|_t|j|_dS)NrrZ)rcrdrLchunk_size_feed_forward seq_len_dimrr$rrlr#ALBERT_ATTENTION_CLASSES_attn_implementationr"rZintermediate_sizer! ffn_outputr hidden_act activationrmrnrorurwrUrVrds  zAlbertLayer.__init__NFrrrroutput_hidden_statesr|cCsL|||||}t|j|j|j|d}|||d}|f|ddS)Nrr)r"rff_chunkrrr#)rvrrrrrrrrUrUrVrszAlbertLayer.forwardrcCs"||}||}||}|Sr)r!rr)rvrrrUrUrVrs   zAlbertLayer.ff_chunkNNFF)rrrrrdrIrrrrrrrrrUrUrwrVrs( rcsteZdZdeffdd Z    ddejdeejdeejd e d e d e e eje ejfd ff d dZ Z S)AlbertLayerGrouprLcs.ttfddtjD|_dS)Ncg|]}tqSrU)r.0rrLrUrV z-AlbertLayerGroup.__init__..)rcrdr ModuleListrangeinner_group_num albert_layersrurwrrVrds $zAlbertLayerGroup.__init__NFrrrrrr|.c Cs|d}d}t|jD]!\}} | |||||} | d}|r#|| df}|r*||f}q |f} |r5| |f} |r<| |f} | S)NrUrr) enumerater) rvrrrrrZlayer_hidden_statesZlayer_attentionsZ layer_indexZ albert_layerZ layer_outputoutputsrUrUrVrs    zAlbertLayerGroup.forwardr)rrrrrdrIrrrrrrrrrUrUrwrVrs&rcsjeZdZdeffdd Z     ddejdeejd eejd e d e d e d e e e ffddZ ZS)AlbertTransformerrLcsFt|_tjj|_tfddt j D|_ dS)NcrrU)rrrrUrVrrz.AlbertTransformer.__init__..) rcrdrLrrrgrembedding_hidden_mapping_inrrnum_hidden_groupsalbert_layer_groupsrurwrrVrds $zAlbertTransformer.__init__NFTrrrrr return_dictr|c Cs||}|r |fnd}|rdnd}|durdg|jjn|}t|jjD]@} t|jj|jj} t| |jj|jj} |j| |||| | | d| ||} | d}|r^|| d}|re||f}q%|sttdd|||fDSt|||dS)NrUrrr\css|] }|dur|VqdSrrU)rvrUrUrV sz,AlbertTransformer.forward..)last_hidden_stater attentions) rrLnum_hidden_layersrrDrrtupler) rvrrrrrrZall_hidden_statesZall_attentionsiZlayers_per_group group_idxZlayer_group_outputrUrUrVrs2    zAlbertTransformer.forward)NNFFT)rrrrrdrIrrrrrrrrrrUrUrwrVrs,  rc@s$eZdZeZeZdZdZddZ dS)AlbertPreTrainedModelalbertTcCst|tjr |jjjd|jjd|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|tj rX|jj |jj ddSt|tre|jj dSdS)zInitialize the weights.r)meanZstdN?) isinstancerrr*rJZnormal_rLZinitializer_ranger-Zzero_rerYr$Zfill_ AlbertMLMHead)rvmodulerUrUrV _init_weights)s       z#AlbertPreTrainedModel._init_weightsN) rrrr config_classrWZload_tf_weightsbase_model_prefixZ_supports_sdparrUrUrUrVr"s  rc@steZdZUdZdZeejed<dZ eejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)AlbertForPreTrainingOutputaQ Output type of [`AlbertForPreTraining`]. 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). sop_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. 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. Nlossprediction_logits sop_logitsrr)rrrrrrrIr__annotations__rrrrrrUrUrUrVr<s rcseZdZeZdZddedeffdd Zdej fdd Z d ej dd fd d Z de e ee fdd fddZe ddeejdeejdeejdeejdeejdeejdeedeedeedeeeffddZZS) AlbertModelrTrLadd_pooling_layercspt|||_t||_t||_|r$t|j |j |_ t |_ nd|_ d|_ |j |_|j|_|dS)zv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer N)rcrdrLrXr(rencoderrrrr&ZTanhpooler_activationrattn_implementationr^ post_init)rvrLrrwrUrVrdcs     zAlbertModel.__init__r|cC|jjSrr(rhrvrUrUrVget_input_embeddingszz AlbertModel.get_input_embeddingsrNcC ||j_dSrr)rvrrUrUrVset_input_embeddings} z AlbertModel.set_input_embeddingsheads_to_prunecCsT|D]#\}}t||jj}t|||jj}|jj|j|j|qdS)a Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} ALBERT has a different architecture in that its layers are shared across groups, which then has inner groups. If an ALBERT model has 12 hidden layers and 2 hidden groups, with two inner groups, there is a total of 4 different layers. These layers are flattened: the indices [0,1] correspond to the two inner groups of the first hidden layer, while [2,3] correspond to the two inner groups of the second hidden layer. Any layer with in index other than [0,1,2,3] will result in an error. See base class PreTrainedModel for more information about head pruning N) itemsrDrLrrrrr"r)rvr layerrrZinner_group_idxrUrUrV _prune_headss  zAlbertModel._prune_headsryrr`r[rrzrrrc Cs|dur|n|jj}|dur|n|jj}| dur| n|jj} |dur*|dur*td|dur9||||} n|durF|dd} ntd| \} } |durU|jn|j} |durctj | | d}|durt |j dr|j j ddd| f}| | | }|}n tj| tj| d}|j ||||d}|jdko|jd ko|duo| }|rt||j| d }n|d d }|j|jd }d|t|jj}|||jj}|j|||||| d}|d}|jdur|||dddfnd}| s||f|d dSt|||j|jdS)NzDYou cannot specify both input_ids and inputs_embeds at the same timer\z5You have to specify either input_ids or inputs_embeds)r~r`r})r[r`rzrr_)Ztgt_lenrr/rar)rrrrr)rZ pooler_outputrr) rLrruse_return_dictrGZ%warn_if_padding_and_no_attention_maskrsr~rIZonesrr(r`rqrrrtrr^rrbZ unsqueezerZfinfominZ get_head_maskrrr&rrrr)rvryrr`r[rrzrrrrrrr~rrZembedding_outputZuse_sdpa_attention_maskZextended_attention_maskZencoder_outputssequence_output pooled_outputrUrUrVrst       *zAlbertModel.forward)T) NNNNNNNNN)rrrrrrrrdrrerr rrDrrrrrIrrrrrrrrUrUrwrVr^sL   rz Albert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `sentence order prediction (classification)` head. )Z custom_introcseZdZddgZdeffdd ZdejfddZd ejdd fd d Z dej fd dZ e dde ejde ejde ejde ejde ejde ejde ejde ejde ede ede edeeeffddZZS)AlbertForPreTrainingpredictions.decoder.biaspredictions.decoder.weightrLcs6t|t||_t||_t||_|dSr) rcrdrrrr' AlbertSOPHeadsop_classifierrrurwrUrVrds     zAlbertForPreTraining.__init__r|cCrrr'decoderrrUrUrVget_output_embeddingsrz*AlbertForPreTraining.get_output_embeddingsnew_embeddingsNcCrrrrvrrUrUrVset_output_embeddingsr z*AlbertForPreTraining.set_output_embeddingscC |jjjSrrr(rhrrUrUrVr z)AlbertForPreTraining.get_input_embeddingsryrr`r[rrzlabelssentence_order_labelrrrc  Cs| dur| n|jj} |j||||||| | | d } | dd\} }|| }||}d}|durU|durUt}||d|jj|d}||dd|d}||}| sl||f| dd}|durj|f|S|St|||| j | j dS)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]` sentence_order_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 original order (sequence A, then sequence B), `1` indicates switched order (sequence B, then sequence A). Example: ```python >>> from transformers import AutoTokenizer, AlbertForPreTraining >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("albert/albert-base-v2") >>> model = AlbertForPreTraining.from_pretrained("albert/albert-base-v2") >>> 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 >>> sop_logits = outputs.sop_logits ```Nrr`r[rrzrrrr/r\)rrrrr) rLrrr'rr rrfrrr)rvryrr`r[rrzr!r"rrrrrrprediction_scoresZ sop_scores total_lossloss_fctmasked_lm_lossZsentence_order_lossoutputrUrUrVrs>)  zAlbertForPreTraining.forward NNNNNNNNNNN)rrr_tied_weights_keysrrdrrrrrerrrrIrrrrrrrrrUrUrwrVrsV      rcsBeZdZdeffdd ZdejdejfddZd d d ZZ S) rrLcspttj|j|jd|_tt|j |_ t |j |j|_ t |j|j |_t|j|_|j |j_ dS)NrZ)rcrdrr$rgrl ParameterrIrrrfr-rrrrr rrrurwrUrVrdVs  zAlbertMLMHead.__init__rr|cCs0||}||}||}||}|}|Sr)rrr$r)rvrr$rUrUrVr`s    zAlbertMLMHead.forwardNcCs,|jjjjdkr|j|j_dS|jj|_dS)Nmeta)rr-r~rrrUrUrV _tie_weightsjszAlbertMLMHead._tie_weights)r|N) rrrrrdrIrrr-rrUrUrwrVrUs  rcs8eZdZdeffdd ZdejdejfddZZS)rrLcs.tt|j|_t|j|j|_ dSr) rcrdrrmclassifier_dropout_probrorr num_labelsr.rurwrUrVrdts zAlbertSOPHead.__init__rr|cCs||}||}|Sr)ror.)rvrZdropout_pooled_outputlogitsrUrUrVrzs  zAlbertSOPHead.forward) rrrrrdrIrrrrUrUrwrVrssrcseZdZddgZfddZdejfddZdejdd fd d Zdej fd d Z e dde e jde e jde e jde e jde e jde e jde e jde ede ede edeeeffddZZS)AlbertForMaskedLMrrcs0t|t|dd|_t||_|dSNF)r)rcrdrrrr'rrurwrUrVrds   zAlbertForMaskedLM.__init__r|cCrrrrrUrUrVrrz'AlbertForMaskedLM.get_output_embeddingsrNcCs||j_|j|j_dSr)r'rr-rrUrUrVrsz'AlbertForMaskedLM.set_output_embeddingscCrrrrrUrUrVrr z&AlbertForMaskedLM.get_input_embeddingsryrr`r[rrzr!rrrc  Cs| dur| n|jj} |j|||||||| | d } | d} || } d}|dur8t}|| d|jj|d}| sN| f| dd}|durL|f|S|St|| | j| j dS)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]` Example: ```python >>> import torch >>> from transformers import AutoTokenizer, AlbertForMaskedLM >>> tokenizer = AutoTokenizer.from_pretrained("albert/albert-base-v2") >>> model = AlbertForMaskedLM.from_pretrained("albert/albert-base-v2") >>> # add mask_token >>> inputs = tokenizer("The capital of [MASK] is Paris.", return_tensors="pt") >>> with torch.no_grad(): ... logits = model(**inputs).logits >>> # retrieve index of [MASK] >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0] >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1) >>> tokenizer.decode(predicted_token_id) 'france' ``` ```python >>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"] >>> labels = torch.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100) >>> outputs = model(**inputs, labels=labels) >>> round(outputs.loss.item(), 2) 0.81 ``` N ryrr`r[rrzrrrrr\r/rr0rr) rLrrr'r rrfrrr)rvryrr`r[rrzr!rrrrZsequence_outputsr$r'r&r(rUrUrVrs61 zAlbertForMaskedLM.forward NNNNNNNNNN)rrrr*rdrrrrrerrrrIrrrrrrrrrUrUrwrVr1sP      r1z Albert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. ceZdZdeffdd Ze          ddeejdeej deejdeejd eej d eej d eejd ee d ee dee de e e ffddZZS)AlbertForSequenceClassificationrLcsRt||j|_||_t||_t|j|_ t |j |jj|_ | dSr)rcrdr/rLrrrrmr.rorrr.rrurwrUrVrds   z(AlbertForSequenceClassification.__init__Nryrr`r[rrzr!rrrr|c  Csr| dur| n|jj} |j|||||||| | d } | d} || } || } d}|dur|jjdurV|jdkr 1` a classification loss is computed (Cross-Entropy). Nr3rZ regressionZsingle_label_classificationZmulti_label_classificationr\r/r4)rLrrror.Z problem_typer/rbrIrtrDr squeezer rr rrr)rvryrr`r[rrzr!rrrrrr0rr&r(rUrUrVrsV     "       z'AlbertForSequenceClassification.forwardr5)rrrrrdrrrIrrrrrrrrrUrUrwrVr7sH     r7cr6)AlbertForTokenClassificationrLcsdt||j|_t|dd|_|jdur|jn|j}t||_ t |j |j j|_ |dSr2)rcrdr/rrr.rnrrmrorrrLr.r)rvrLr.rwrUrVrdIs    z%AlbertForTokenClassification.__init__Nryrr`r[rrzr!rrrr|c  Cs| dur| n|jj} |j|||||||| | d } | d} || } || } d}|durrsclampr rrr)rvryrr`r[rrzr<r=rrrrrr0r>r?r%Z ignored_indexr&Z start_lossZend_lossr(rUrUrVrsP          z"AlbertForQuestionAnswering.forwardr)rrrrrdrrrIrrrrrrrrrUrUrwrVr:sN      r:cr6)AlbertForMultipleChoicerLcs@t|t||_t|j|_t|j d|_ | dS)Nr) rcrdrrrrmr.rorrr.rrurwrUrVrds   z AlbertForMultipleChoice.__init__Nryrr`r[rrzr!rrrr|c  Csn| dur| n|jj} |dur|jdn|jd} |dur%|d|dnd}|dur4|d|dnd}|durC|d|dnd}|durR|d|dnd}|dure|d|d|dnd}|j|||||||| | d } | d} || } || }|d| }d}|durt}|||}| s|f| dd}|dur|f|S|St ||| j | 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.__call__`] and [`PreTrainedTokenizer.encode`] 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) Nrr\rr#r/r4) rLrrFrrsrror.r rrr)rvryrr`r[rrzr!rrrZ num_choicesrrr0Zreshaped_logitsrr&r(rUrUrVrsL,    zAlbertForMultipleChoice.forwardr5rArUrUrwrVrBsH     rB) rWrrrr1r7r9r:rB)Brrr5 dataclassesrtypingrrrrrrIrZtorch.nnr r r Z activationsr Zmodeling_attn_mask_utilsrZmodeling_outputsrrrrrrrZmodeling_utilsrZ pytorch_utilsrrrrutilsrrrZconfiguration_albertrZ get_loggerrr3rWModulerXrrrrrrrrrrrrr1r7r9r:rB__all__rUrUrUrVsh    $   ~Dp7)#5! f iWGM f