o Zh@sdZddlZddlmZmZmZddlZddlZddlmZddl m Z m Z m Z m Z ddl mZddlmZdd lmZdd lmZmZmZmZmZdd lmZdd lmZmZd dlm Z e!e"Z#d(ddZ$Gdddej%Z&Gdddej%Z'Gdddej%Z(eGdddeZ)eGddde)Z*eddGddde)eZ+ed dGd!d"d"e)Z,eGd#d$d$e)Z-eGd%d&d&e)Z.gd'Z/dS))zPyTorch MPT model.N)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLoss LayerNormMSELoss) functional)GenerationMixin)!_prepare_4d_causal_attention_mask))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentionsQuestionAnsweringModelOutput SequenceClassifierOutputWithPastTokenClassifierOutput)PreTrainedModel)auto_docstringlogging) MptConfigcCstjd|dtj|dddd|}dtt|}tjd|dtj|d}|||}dt d|}|d|dd}||krjtj |ddddddf|ddddddfgddddd|df}||}| dS) a Link to paper: https://arxiv.org/abs/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_2baseZslopesr/S/var/www/auris/lib/python3.10/site-packages/transformers/models/mpt/modeling_mpt.pybuild_mpt_alibi_tensor+s$ L r1c sZeZdZdZdeffdd Z  d dejdejdee ejd eejfd d Z Z S) MptAttentionzzMulti-head self attention. Using torch or triton attention implementation enables user to also use additive bias. configcst|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_projselfr3 __class__r/r0r7Gs     zMptAttention.__init__N hidden_states position_biaspast_key_valueattention_maskcCs4|jdd\}}||}|jr|j|j |jd}|jddd\}} } ||||j|jdd}| |||j|jdd} | |||j|jdd} |durtt |dkrot j |d| gdd} t j |d| gdd} | | f}n| | f}t || dd|j } |dur|n||djd} |durt |jdkrtd t |j| jd} td|d| }td|d| }|dd|d|df}| |} |dur| |t |jj} tjj| dd| j}tjj||j|jd }t || }|dddd||d}| |}|||fS) Nr)minmaxr rrrz6Expecting position_bias shape to be 3 dimensions, got ptraining)!shaperCrAclampchunkreshaper9r;Z transposelenrcatmatmulr= ValueErrorrNsizeZ masked_fillZfinforrMrr Zsoftmaxr&todropoutr@rSZpermute contiguousr"rD)rFrIrJrKrL batch_size seq_lengthZ mixed_qkvZ query_statesZ key_statesZ value_statesZattention_scoresZ query_lengthZ key_lengthZposition_bias_query_indexZposition_bias_key_index attn_weightsZcontext_statesZ attn_outputr/r/r0forwardVs@       zMptAttention.forward)NN) __name__ __module__ __qualname____doc__rr7rTensorrrrc __classcell__r/r/rGr0r2Bs r2cs>eZdZdeffdd ZdejdejdejfddZZS) MptMLPr3csXt|j}tj|d|dd|_tjdd|_tjd||dd|_|j j |_ dS)NFr4none)Z approximate) r6r7r8rrBup_projZGELUact down_projr<r?hidden_dropoutrFr3r8rGr/r0r7s zMptMLP.__init__rIresidualreturncCs:|||}||}tj||j|jd}||}|S)NrQ)rnrmroFr^rprS)rFrIrrZintermediate_outputoutputr/r/r0rcs  zMptMLP.forward) rdrerfrr7rrhrcrir/r/rGr0rjs$ rjc sbeZdZdeffdd Z   ddejdejdejd eeejejfd e d e f d d Z Z S)MptBlockr3csxt|j}t||jd|_d|j_|j|_t ||_ t||jd|_ d|j _t ||_ |jj|_t|j|_dS)NZeps)r6r7r8rlayer_norm_epsilonnorm_1r5r9r*r2attnnorm_2rjffnr<r?Z dropout_raterDropoutresid_attn_dropoutrqrGr/r0r7s    zMptBlock.__init__NFrIrJrL layer_past use_cacheoutput_attentionscCst||}|}|j||||d\} } } || |}||}|}|||} | f} |r1| | f7} |r8| | f7} | S)N)rJrLrK)ryrzr~r{r|)rFrIrJrLrrrZlayernorm_outputrrZ attn_outputsrbrKruoutputsr/r/r0rcs$      zMptBlock.forward)NFF) rdrerfrr7rrhrrboolrcrir/r/rGr0rvs$rvcszeZdZeZdZdZdgZdgZfddZ de j fdd Z e d eeejejfd eeejejffd d ZZS)MptPreTrainedModel transformerTrvz lm_head.*.cstj|i|dSN)r6r7)rFZinputskwargsrGr/r0r7szMptPreTrainedModel.__init__modulecCst|tjr |jjjd|jjd|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|t r\|jdurS|jj |jj ddSdS)zInitialize the weights.g)meanZstdNr) isinstancerrBweightdataZnormal_r3Zinitializer_ranger5Zzero_ EmbeddingZ padding_idxrZfill_)rFrr/r/r0 _init_weightss       z MptPreTrainedModel._init_weightsrKrscs8|ddj\}}||tfdd|DS)zw Converts the cache to the format expected by Mpt, i.e. to tuple(tuple([batch_size * num_heads, ...])) rc3s4|]}|d|dfVqdSrrN)rW.0rZbatch_size_times_num_headsr;rar/r0 s z;MptPreTrainedModel._convert_to_mpt_cache..)rTtuple)rKr`r*r/rr0_convert_to_mpt_caches z(MptPreTrainedModel._convert_to_mpt_cache)rdrerfrZ config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_keys_to_ignore_on_load_missingr7rModuler staticmethodrrrhrrir/r/rGr0rs rcseZdZdeffdd ZddZddd Zd ejfd d Z e        dd e ej de e e ejejfdfde ejde ej de ede ede ede edee ejdfeffddZZS)MptModelr3csztj|_j|_tj|j|_t fddt j D|_ t |jjd|_d|j_d|_|dS)Ncsg|]}tqSr/)rv)r_r3r/r0 sz%MptModel.__init__..rwF)r6r7r8r9r*rr vocab_sizewteZ ModuleListrangeZn_layersblocksrrxnorm_fr5gradient_checkpointing post_initrErGrr0r7s   zMptModel.__init__cC|jSrrrFr/r/r0get_input_embeddings+zMptModel.get_input_embeddingsrNcCst||||Sr)r1)rFr*r+r,rr/r/r0r1.szMptModel.build_mpt_alibi_tensornew_embeddingscC ||_dSrrrFrr/r/r0set_input_embeddings1 zMptModel.set_input_embeddings input_idspast_key_values.rL inputs_embedsrroutput_hidden_states return_dictrsc  Ks~|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur$|n|jj}|dur4|dur4td|dur>|j\} } n|durI|j\} } } ntd|dur[tdgt|j }|durd| |}|} |rjdnd}|rpdnd}|rvdnd}|j r|j r|rt dd}| }d}|ddur|ddjd}||}|durtj| |f| jd }n|| j}|j|j|jj| jd }t|| | f||}|}t|j |D]G\}}|r|| f}|j r|j r||j| |||||}n || |||||d }|d} |d ur ||d f}|r|||rdnd f}q|| } |r'|| f}|s7td 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[0][0].shape[2]` (`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_embedsr/zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Frrr)rrLrrrJTrcss|] }|dur|VqdSrr/)rvr/r/r0rsz#MptModel.forward..)Zlast_hidden_staterrI attentions)r3rrruse_return_dictr[rTrrXrrrrSlogger warning_oncerZonesrr]r1r*r:r rzipZ_gradient_checkpointing_func__call__rr)rFrrrLrrrrrrr`rarrIZpresentsZall_self_attentionsZall_hidden_statesZseq_length_with_pastZpast_key_values_lengthr-Z causal_maskblockrrr/r/r0rc4s               zMptModel.forwardrNNNNNNNNN)rdrerfrr7rr1rrhrrr LongTensorrrrrrcrir/r/rGr0rsB   rz The MPT Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). )Z custom_introcseZdZdgZdeffdd ZddZdejfdd Z e dd e ej d e e e ejejfd fde ejde ejde ejde ede ede ede edee ejeffddZde e ejejfd fdej de e ejejfd ffddZZS)MptForCausalLMzlm_head.weightr3cs8t|t||_tj|j|jdd|_| dSNFr4) r6r7rrrrBr8rlm_headrrErGr/r0r7s   zMptForCausalLM.__init__cCrrrrr/r/r0get_output_embeddingsrz$MptForCausalLM.get_output_embeddingsrcCrrrrr/r/r0set_output_embeddingsrz$MptForCausalLM.set_output_embeddingsNrr.rLrlabelsrrrrrsc  Ks| dur| n|jj} |j|||||||| d} | d} || } d}|dur:|| j}|j| |fd|jji| }| sP| f| dd}|durN|f|S|St|| | j | j | j dS)aZ 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[0][0].shape[2]` (`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]` NrrLrrrrrrrrlosslogitsrrIr) r3rrrr]rZ loss_functionrrrrIr)rFrrrLrrrrrrrtransformer_outputsrIZ lm_logitsrrur/r/r0rcsD  zMptForCausalLM.forwardpastbeam_idxcs,fdd|Dtfdd|D}|S)aL This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct beam_idx at every generation step. Output shares the same memory storage as `past`. cs&i|]}|D] }|j|jqqSr/)rr])rrZ past_state)rr/r0 s z1MptForCausalLM._reorder_cache..c3sD|]}|dd|dj|dd|djfVqdSr)Z index_selectrr)device_to_beam_idxr/r0rs z0MptForCausalLM._reorder_cache..)r)rFrrZreordered_pastr/)rrr0_reorder_caches  zMptForCausalLM._reorder_cache NNNNNNNNN)rdrerfZ_tied_weights_keysrr7rrrhrrrrrrrrrcrrir/r/rGr0rsV   Fra 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). ceZdZdeffdd Ze         ddeejdee e ej ej fdfdeej d eej d eej d ee d ee d ee dee de e ej e ffddZZS)MptForSequenceClassificationr3cs@t||j|_t||_tj|j|jdd|_| dSr) r6r7 num_labelsrrrrBr8scorerrErGr/r0r78s   z%MptForSequenceClassification.__init__Nrr.rLrrrrrrrsc  Cs$| dur| n|jj} |j|||||||| d} | d} || } |dur*|jd} n|jd} |jjdur=| dkr=td|jjdurFd}n1|durk||jjk| jt j }t j |jd| jt j d}|| d}n d}t |jjd| t j | | jd |f}d}|dur|jjdur|jdkrd |j_n|jdkr|jt jks|jt jkrd |j_nd |j_|jjd krt}|jdkr|||}n#|||}n|jjd krt}|||}n|jjd krt}|||}| s|f| dd}|dur|f|S|St||| j| j| jd S)4 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[0][0].shape[2]` (`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). Nrrrz=Cannot handle batch sizes > 1 if no padding token is defined.rO)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)r3rrrrTZ pad_token_idr[r]rrr!r ZargmaxrrrHrdZ problem_typerrlongintr r)rrrrrIr)rFrrrLrrrrrrrrIrr`Zlast_non_pad_tokenZ non_pad_maskZ token_indicesZ pooled_logitsrloss_fctrur/r/r0rcAst        "        z$MptForSequenceClassification.forwardr)rdrerfrr7rrrrrrhrrrrcrir/r/rGr0r)sB    rcr)MptForTokenClassificationr3cst||j|_t||_t|dr|jdur|j}nt|dr+|jdur+|j}nd}t ||_ t |j |j|_ |dS)Nclassifier_dropoutrpg?)r6r7rrrhasattrrrprr}r^rBr8 classifierr)rFr3rrGr/r0r7s    z"MptForTokenClassification.__init__Nrr.rLrrrrrrrsc  Ks| dur| n|jj} |j|||||||| d} | d} || } || } d}|durJ|| j}|j\}}t}|| |||j | ||}| s`| f| dd}|dur^|f|S|St || | j | j dS)rNrrr)rrrIr)r3rrr^rr]rrTrr"rrrIr)rFrrrLrrrrrrZdeprecated_argumentsrrIrrr`rarrur/r/r0rcs>    z!MptForTokenClassification.forwardr)rdrerfrr7rrrrrrhrrrrcrir/r/rGr0rsB   rcseZdZfddZe        ddeejdeejdeejdeejdeejd ee d ee d ee d e e e ffd dZ ZS)MptForQuestionAnsweringcs2t|t||_t|jd|_|dS)Nr) r6r7rrrrBr8 qa_outputsrrErGr/r0r7s   z MptForQuestionAnswering.__init__NrrLrstart_positions end_positionsrrrrsc CsB|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 | j d S) rN)rLrrrrrrrOr)Z ignore_indexr)r start_logits end_logitsrIr)r3rrrsplitr)r_rXr\rUrrrIr)rFrrLrrrrrrrZsequence_outputrrrZ total_lossZ ignored_indexrZ start_lossZend_lossrur/r/r0rc sJ          zMptForQuestionAnswering.forwardr)rdrerfr7rrrrZ FloatTensorrrrrrcrir/r/rGr0rs<   r)rrrrrrr)0rgr#typingrrrrZtorch.utils.checkpointrZtorch.nnrrrr r rtZ generationr Zmodeling_attn_mask_utilsr Zmodeling_outputsrrrrrZmodeling_utilsrutilsrrZconfiguration_mptrZ get_loggerrdrr1rr2rjrvrrrrrr__all__r/r/r/r0sJ        L@/ prX R