a h.@sdZddlZddlmZmZddlZddlZddlmZddlm Z m Z m Z m Z ddlm ZddlmZmZdd lmZdd lmZdd lmZdd lmZmZmZmZmZdd lmZddl m!Z!m"Z"ddl#m$Z$ddl%m&Z&e"'e(Z)d+ddZ*Gdddej+Z,Gdddej+Z-GdddeZ.e!GdddeZ/e!Gddde/Z0e!dd Gd!d"d"e/eZ1e!d#d Gd$d%d%e/Z2e!Gd&d'd'e/Z3e!Gd(d)d)e/Z4gd*Z5dS),zPyTorch MPT model.N)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLoss LayerNormMSELoss) functional)Cache DynamicCache)GenerationMixin)!_prepare_4d_causal_attention_mask)GradientCheckpointingLayer))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentionsQuestionAnsweringModelOutput SequenceClassifierOutputWithPastTokenClassifierOutput)PreTrainedModel)auto_docstringlogging)deprecate_kwarg) MptConfigcCstjd|dtj|dddd|}dtt|}tjd|dtj|d}|||}dt d|}|d|dd}||krtj |ddddddf|ddddddfgddddd|df}||}| dS) a Link to paper: https://huggingface.co/papers/2108.12409 - Alibi tensor is not causal as the original paper mentions, it relies on a translation invariance of softmax for quick implementation. This implementation has been copied from the alibi implementation of MPT source code that led to slightly different results than the Bloom alibi: https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L292 r)dtypedevice?N.dimr) torcharangeint32viewmathceillog2Zint64floatpowconcatsqueeze) num_headssequence_lengthalibi_bias_maxralibiZnum_heads_power_of_2baseZslopesr2`/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/mpt/modeling_mpt.pybuild_mpt_alibi_tensor.s$ Lr4csheZdZdZd eeedfdd Zedddd de j e j ee ee j ee j d d d Z Z S) MptAttentionzzMulti-head self attention. Using torch or triton attention implementation enables user to also use additive bias. Nconfig layer_idxcst|j|_|j|_|j|_|j|j|_|jj|_|jdur\dt |j|j|_|jj |_ |jj |_ tj|jd|jdd|_tj|j|jdd|_||_dS)Nrr Fbias)super__init__ hidden_sizen_heads max_seq_lenZmax_seq_lengthhead_dim attn_config softmax_scaler&sqrt attn_pdropattn_dropout_pclip_qkvrLinearWqkvout_projr8)selfr7r8 __class__r2r3r<Js     zMptAttention.__init__past_key_valuepast_key_values4.58new_nameversion) hidden_states position_biasrNattention_maskcache_positioncCs|jdd\}}||}|jr6|j|j |jd}|jddd\} } } | |||j|jdd} | |||j|jdd} | |||j|jdd} |durd|i} | | | |j | \} } t | | dd|j } |dur|n ||}|durxt|jdkr$td t|j| jd}td |d|}td |d|}|dd|d|df}| |} |dur| |t | jj} tjj| dd| j}tjj||j|jd }t || }|d ddd !||d}|"|}||fS) Nr)minmaxr r rrVz6Expecting position_bias shape to be 3 dimensions, got rptraining)#shaperHrFclampchunkreshaper>r@Z transposeupdater8r"matmulrBget_seq_lengthlen ValueErrorrXsizeZ masked_fillZfinforrWrr Zsoftmaxr)todropoutrEr]Zpermute contiguousr%rI)rJrSrTrNrUrV batch_size seq_lengthZ mixed_qkvZ query_statesZ key_statesZ value_statesZ cache_kwargsZattention_scoresZ query_lengthZ key_lengthZposition_bias_query_indexZposition_bias_key_index attn_weightsZcontext_statesZ attn_outputr2r2r3forwardZs:       zMptAttention.forward)N)NNN)__name__ __module__ __qualname____doc__rrintr<rr"Tensorr rn __classcell__r2r2rKr3r5Es r5cs:eZdZedfdd ZejejejdddZZS)MptMLPr7csXt|j}tj|d|dd|_tjdd|_tjd||dd|_|j j |_ dS)NFr9none)Z approximate) r;r<r=rrGup_projZGELUact down_projrArDhidden_dropout)rJr7r=rKr2r3r<s  zMptMLP.__init__)rSresidualreturncCs:|||}||}tj||j|jd}||}|S)Nr[)r{rzr|Frir}r])rJrSr~Zintermediate_outputoutputr2r2r3rns  zMptMLP.forward) rorprqrr<r"rtrnrur2r2rKr3rvs rvc sVeZdZd eeedfdd Zd ejejejee e e eejdddZ Z S) MptBlockNr6cszt|j}t||jd|_d|j_|j|_t |||_ t||jd|_ d|j _t ||_ |jj|_t|j|_dS)Neps)r;r<r=rlayer_norm_epsilonnorm_1r:r>r-r5attnnorm_2rvffnrArDZ dropout_raterDropoutresid_attn_dropout)rJr7r8r=rKr2r3r<s    zMptBlock.__init__F)rSrTrU layer_past use_cacheoutput_attentionsrVc CsV||}|} |j|||||d\} } || | }||}|} ||| } | | fS)N)rTrUrNrV)rrrrr) rJrSrTrUrrrrVZlayernorm_outputr~Z attn_outputsrmrr2r2r3rns    zMptBlock.forward)N)NFFN) rorprqrrrsr<r"rtr boolrnrur2r2rKr3rsrcseZdZUeed<dZdZdgZdgZfddZ e j dd d Z e ed d d deeejejfeeejejfdddZZS)MptPreTrainedModelr7 transformerTrz lm_head.*.cstj|i|dSN)r;r<)rJinputskwargsrKr2r3r<szMptPreTrainedModel.__init__)modulecCst|tjr:|jjjd|jjd|jdur|jj nnt|tj rz|jjjd|jjd|j dur|jj|j  n.t|t r|jdur|jj |jj ddS)zInitialize the weights.g)meanZstdNr) isinstancerrGweightdataZnormal_r7Zinitializer_ranger:Zzero_ EmbeddingZ padding_idxrZfill_)rJrr2r2r3 _init_weightss       z MptPreTrainedModel._init_weightsrMrNrOrP)rNrcs8|ddj\}}||tfdd|DS)zw Converts the cache to the format expected by Mpt, i.e. to tuple(tuple([batch_size * num_heads, ...])) rc3s2|]*}|d|dfVqdS)rrN)ra).0rZbatch_size_times_num_headsr@rlr2r3 sz;MptPreTrainedModel._convert_to_mpt_cache..)r^tuple)rNrkr-r2rr3_convert_to_mpt_caches z(MptPreTrainedModel._convert_to_mpt_cache)rorprqr__annotations__Zbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_keys_to_ignore_on_load_missingr<rModuler staticmethodrrr"rtrrur2r2rKr3rs   rcseZdZedfdd ZddZddd Zejd d d Z e de ej e e eeejejfd fefe eje ej e ee ee ee ee eje eejd fefd ddZZS)MptModelrwcsztj|_j|_tj|j|_t fddt j D|_ t |jjd|_d|j_d|_|dS)Ncsg|]}t|dqS))r8)r)rirwr2r3 z%MptModel.__init__..rF)r;r<r=r>r-rr vocab_sizewteZ ModuleListrangeZn_layersblocksrrnorm_fr:gradient_checkpointing post_initrJr7rKrwr3r<s  zMptModel.__init__cCs|jSrr)rJr2r2r3get_input_embeddings'szMptModel.get_input_embeddingsrNcCst||||Sr)r4)rJr-r.r/rr2r2r3r4*szMptModel.build_mpt_alibi_tensornew_embeddingscCs ||_dSrrrJrr2r2r3set_input_embeddings-szMptModel.set_input_embeddings.) input_idsrNrU inputs_embedsrroutput_hidden_states return_dictrVrc  Ks@|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durH|n|jj}|durj|durjtdn2|dur~|j\} } n|dur|j\} } } ntd|jr|jr|rt dd}|dur| |}|r|durt |jd}|r t |tr t dt |}|}|rdnd}|r&dnd}|dur<|nd }| |}|durhtj| |f|jd }n ||j}|j|j|jj|jd }t|| | f||}|}|jD]H}|r||f}|||||||| d }|d }|r||d f}q||}|r||f}|s0td d||||fDSt||||dS) input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) NzDYou cannot specify both input_ids and inputs_embeds at the same timez5You have to specify either input_ids or inputs_embedszZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...FrwzPassing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. You should pass an instance of `DynamicCache` instead, e.g. `past_key_values=DynamicCache.from_legacy_cache(past_key_values)`.r2rr)rrUrrrTrVrcss|]}|dur|VqdSrr2)rvr2r2r3rsz#MptModel.forward..)Zlast_hidden_staterNrS attentions)r7rrruse_return_dictrfr^rr]logger warning_oncerr rrZfrom_legacy_cacherdr"Zonesrrhr4r-r?rrrrr)rJrrNrUrrrrrrVrrkrl_rSZall_self_attentionsZall_hidden_statesZpast_key_values_lengthZseq_length_with_pastr0Z causal_maskblockoutputsr2r2r3rn0s                zMptModel.forward)rN) NNNNNNNNN)rorprqrr<rr4r"rtrrr LongTensorrrr rrrnrur2r2rKr3rs4 "rz The MPT Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). )Z custom_introcseZdZdgZedfdd ZejdddZe d e ej e e e ejejfd fe eje eje eje e e e e e e e e ejee ejefd d d ZZS)MptForCausalLMzlm_head.weightrwcs8t|t||_tj|j|jdd|_| dSNFr9) r;r<rrrrGr=rlm_headrrrKr2r3r<s  zMptForCausalLM.__init__rcCs ||_dSr)rrr2r2r3set_output_embeddingssz$MptForCausalLM.set_output_embeddingsN.) rrNrUrlabelsrrrrrVrc  Ks| dur | n|jj} |j|||||||| | d } | d} || }d}|durv||j}|j||fd|jji| }| s|f| dd}|dur|f|S|St||| j | j | j dS)a\ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` N)rNrUrrrrrrVrrrlosslogitsrNrSr) r7rrrrhrZ loss_functionrrrNrSr)rJrrNrUrrrrrrrVrtransformer_outputsrSZ lm_logitsrrr2r2r3rnsF    zMptForCausalLM.forward) NNNNNNNNNN)rorprqZ_tied_weights_keysrr<r"rtrrrrrrrrrnrur2r2rKr3rs6ra The MPT Model transformer with a sequence classification head on top (linear layer). [`MptForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-1) do. Since it does classification on the last token, it requires to know the position of the last token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). cseZdZedfdd Zed eejee e ej ej fdfeej eej eej ee ee ee ee e e ej e fd ddZZS) MptForSequenceClassificationrwcs@t||j|_t||_tj|j|jdd|_| dSr) r;r< num_labelsrrrrGr=scorerrrKr2r3r<s   z%MptForSequenceClassification.__init__N. rrNrUrrrrrrrc  Cs8| dur | n|jj} |j|||||||| d} | d} || } |durT|jd} n |jd} |jjdurz| dkrztd|jjdurd}nb|dur||jjk| jt j }t j |jd| jt j d}|| d}nd}t |jjd| t j | | jd |f}d}|dur|jjdurr|jdkr8d |j_n:|jdkrj|jt jks`|jt jkrjd |j_nd |j_|jjd krt}|jdkr|||}n |||}n>|jjd krt}|||}n|jjd krt}|||}| s |f| dd}|dur|f|S|St||| j| j| jd S)6 input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) 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). NrNrUrrrrrrrz=Cannot handle batch sizes > 1 if no padding token is defined.rY)rrz will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`rZ regressionZsingle_label_classificationZmulti_label_classificationr)r7rrrr^Z pad_token_idrfrhrr"r$r#ZargmaxrrrLroZ problem_typerrlongrsrr,rrrrNrSr)rJrrNrUrrrrrrrrSrrkZlast_non_pad_tokenZ non_pad_maskZ token_indicesZ pooled_logitsrloss_fctrr2r2r3rnst         (     z$MptForSequenceClassification.forward) NNNNNNNNN)rorprqrr<rrr"rrrtrrrrnrur2r2rKr3rs. rcseZdZedfdd Zed eejee e ej ej fdfeej eej eej ee ee ee ee e e ej e fd ddZZS) MptForTokenClassificationrwcst||j|_t||_t|dr:|jdur:|j}n t|drV|jdurV|j}nd}t ||_ t |j |j|_ |dS)Nclassifier_dropoutr}g?)r;r<rrrhasattrrr}rrrirGr= classifierr)rJr7rrKr2r3r<s   z"MptForTokenClassification.__init__N.rc  Ks| dur | n|jj} |j|||||||| d} | d} || } || } d}|dur|| j}|j\}}t}|| |||j | ||}| s| f| dd}|dur|f|S|St || | j | j dS)rNrrr)rrrSr)r7rrrirrhrr^rr%rrrSr)rJrrNrUrrrrrrZdeprecated_argumentsrrSrrrkrlrrr2r2r3rns>     z!MptForTokenClassification.forward) NNNNNNNNN)rorprqrr<rrr"rrrtrrrrnrur2r2rKr3rs.rc sreZdZfddZedeejeejeejeejeejee ee ee e e e fd ddZ ZS)MptForQuestionAnsweringcs2t|t||_t|jd|_|dS)Nr) r;r<rrrrGr= qa_outputsrrrKr2r3r<s  z MptForQuestionAnswering.__init__N) rrUrstart_positions end_positionsrrrrc CsF|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 | j d S) rN)rUrrrrrrrYr )Z ignore_indexr)r start_logits end_logitsrSr)r7rrrsplitr,rjrergr_rrrSr)rJrrUrrrrrrrZsequence_outputrrrZ total_lossZ ignored_indexrZ start_lossZend_lossrr2r2r3rnsJ           zMptForQuestionAnswering.forward)NNNNNNNN)rorprqr<rrr"rZ FloatTensorrrrrrnrur2r2rKr3rs*  r)rrrrrr)rN)6rrr&typingrrr"Ztorch.utils.checkpointrZtorch.nnrrrrr rZ cache_utilsr r Z generationr Zmodeling_attn_mask_utilsrZmodeling_layersrZmodeling_outputsrrrrrZmodeling_utilsrutilsrrZutils.deprecationrZconfiguration_mptrZ get_loggerrorr4rr5rvrrrrrrr__all__r2r2r2r3sN          J:0XrXR