T^(h=ddlmZmZmZmZmZddlZddlZddlmZddl m Z m Z m Z ddl mZddlmZddlmZmZdd lmZdd lmZdd lmZd d lmZmZmZmZmZm Z m!Z!m"Z"m#Z#m$Z$ddl%m&Z&ejNe(Z)dZ*Gdde!Z+GddeZ,GddeZ-Gdde"Z.GddeZ/Gdde Z0GddeZ1Gdd eZ2y)!)CallableListOptionalTupleUnionN)nn)CacheSlidingWindowCache StaticCache)AttentionMaskConverter)FlashAttentionKwargs)BaseModelOutputWithPastQuestionAnsweringModelOutput)ALL_ATTENTION_FUNCTIONS)Unpack)logging) LlamaAttentionLlamaDecoderLayerLlamaForCausalLMLlamaForQuestionAnsweringLlamaForSequenceClassificationLlamaForTokenClassificationLlamaMLP LlamaModelapply_rotary_pos_embeager_attention_forward) MistralConfigzmistralai/Mistral-7B-v0.1ceZdZfdZxZS) MistralMLPcJt||tj|j|j d|_tj|j|j d|_tj|j |jd|_yNF)bias) super__init__rLinear hidden_sizeintermediate_size gate_projup_proj down_projselfconfig __class__s i/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/transformers/models/mistral/modular_mistral.pyr'zMistralMLP.__init__#ss  4#3#3T5K5KRWXyy!1!143I3IPUV 4#9#94;K;KRWX)__name__ __module__ __qualname__r' __classcell__r1s@r2r"r""sYYr3r"c2eZdZdedeffd Z d dejdeejejfde ejde e de ejd e e d eeje eje eejffd ZxZS) MistralAttentionr0 layer_idxct|tj|j|j |j zd|_tj|j|j|j zd|_ tj|j|j|j zd|_ tj|j |j z|jd|_ yr$) r&r'rr(r)num_attention_headshead_dimq_projnum_key_value_headsk_projv_projo_projr/r0r;r1s r2r'zMistralAttention.__init__+s ii 2 2F4N4NQUQ^Q^4^ejk ii 2 2F4N4NQUQ^Q^4^ejk ii 2 2F4N4NQUQ^Q^4^ejk ii : :T]] JFL^L^ejk r3 hidden_statesposition_embeddingsattention_maskpast_key_valuecache_positionkwargsreturnc |jdd}g|d|j}|j|j|j dd} |j |j|j dd} |j |j|j dd} |\} } t| | | | \} } |'| | |d}|j| | |j|\} } t}|jjdk7r^|jjdk(r(|jddrtjd nt |jj}||| | | |f|j"sd n |j$|j&t)|jd dd |\}}|j*g|dj-}|j/|}||fS) Nrr)sincosrIeagersdpaoutput_attentionsFz`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.sliding_window)dropoutscalingrT)shaper>r?view transposerArBrupdater;rr0_attn_implementationgetlogger warning_oncertrainingattention_dropoutrVgetattrreshape contiguousrC)r/rErFrGrHrIrJ input_shape hidden_shape query_states key_states value_statesrOrN cache_kwargsattention_interface attn_output attn_weightss r2forwardzMistralAttention.forward2s$))#2. 88b8$--8 {{=166|DNNqRST [[/44\BLLQPQR {{=166|DNNqRST &S#7 jRUWZ#[ j  %#&snUL'5'<'d##L '>dkk>^>^&_#$7      % $}}C$2H2HLL"4;;0@$G %  % ! \*k));;;;FFH kk+. L((r3)NN)r4r5r6r intr'torchTensorrrr LongTensorrrrmr7r8s@r2r:r:*sl}ll+/59 0)||0)#5<<#=>0)!. 0) ! 0) !!1!12 0)-.0) u||Xell3XeELL>Q5RR S0)r3r:c(eZdZdedeffd ZxZS)MistralDecoderLayerr0r;cjt|||t|||_t ||_y)N)r0r;)r&r'r: self_attnr"mlprDs r2r'zMistralDecoderLayer.__init__fs, +)9Mf%r3)r4r5r6r rnr'r7r8s@r2rsrses&}&&&r3rsceZdZdeffd Z ddej dej dej dedef dZ e dej d e d e d ejd ejdej d e dedefdZxZS) MistralModelr0c t||tjt |j Dcgc]}t ||c}|_ycc}wN)r&r'r ModuleListrangenum_hidden_layersrslayersrDs r2r'zMistralModel.__init__msD  mmEJ6KcKcEd e  3 e  esArG input_tensorrIpast_key_valuesrRc |jjdk(rS|H|F|dddfjj|j dk7}|r t d|d|vr|Sy||j nd}t|t}t|t} |jjdk(r?|s=| s;|s9tj||||jj|jry|j|j} } t!j"| j$} |j&d} | s|r|j)}n1t|t j*r|j&dn|| zdz}|j-|| || | ||j&d|j| }|jjdk(r2|0|jj.d vr|stj0|| }|S) Nflash_attention_2rMrzYou are attempting to perform batched generation with padding_side='right' this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to call `tokenizer.padding_side = 'left'` before tokenizing the input. rSrQ) inputs_embedspast_key_values_lengthrT is_trainingr)sequence_length target_lengthdtypedevicerI batch_sizer0r)cudaxpu)r0r[sumitemsize ValueErrorget_seq_length isinstancer r r _ignore_causal_mask_sdparTr_rrrofinfominrWget_max_cache_shaperp5_prepare_4d_causal_attention_mask_with_cache_positiontype_unmask_unattended)r/rGrrIrrRis_padding_rightpast_seen_tokensusing_static_cacheusing_sliding_window_cacherr min_dtyperr causal_masks r2_update_causal_maskz MistralModel._update_causal_maskss ;; + +/B B)o.I#1!R%#8#<#<#>#C#C#EIZIZI\]^I_#_ #$a )c^.C%% @O?Z?99;`a'E%/AS%T" KK , , 6'+E%%>>*'7#{{99 MM $**L,?,?vKK&** &,,Q/ %);+??AM nell;$$R(%7!; PP +')#))!,;;+Q   KK , , 6*%%**o=% 1CCKQZ[Kr3rrrrrc p||jdk(r|} | Stj|j} tj||f| ||} tj |||j ddkD} |j]t|tr||kDrHtj |||j dd|jz k} | j| | | z} | ddddddfj|ddd} || j} |jd|kDr |ddd|f}|jd} | ddddddd| f|ddddddfj| jz}|dk(}| ddddddd| fj!|| | ddddddd| f<| S)aV Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. device (`torch.device`): The device to place the 4D attention mask on. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. config (`MistralConfig`): The model's configuration class past_key_values (`Cache`): The cache class that is being used currently to generate N) fill_valuerr)rrMrr)dimrorrfullarangerbrTrr bitwise_or_expandclonerWtor masked_fill)rGrrrrrIrr0rrrdiagonal_attend_masksliding_attend_mask mask_length padding_masks r2rzBMistralModel._prepare_4d_causal_attention_mask_with_cache_positionsH  %.*<*<*>!*C(K:7 E*..I** -0Ye\bK$)<< f#MP^PfPfgiklPm#m $$0"/3EF/\iJi*/,,}V*T&..r158M8MM+')445HI / /K%dD!Q&67>>z1bRTUK))//1 !''+m;%3A~ ~4E%FN,2226 *1aL[L+@ANSTVZ\`bcScDdDgDg&&E   ,q0 5@Aq,;,AV5W5c5c )6 Aq!\k\12r3)F)r4r5r6r r'rorpr boolr staticmethodrnrrrr7r8s@r2rxrxls } #( Q QllQ Q  Q  QfB BBB{{ B  B  BBBBBr3rxc eZdZy)MistralForCausalLMNr4r5r6r3r2rr r3rc eZdZy)MistralForTokenClassificationNrrr3r2rrrr3rc eZdZy) MistralForSequenceClassificationNrrr3r2rrrr3rcDeZdZdZfdZdZdZ ddeejdeejdeejdee e e ejfd eejd eejd eejd eed eedefdZxZS)MistralForQuestionAnsweringmodelcHt||t||_|`yrz)r&r'rxr transformerr.s r2r'z$MistralForQuestionAnswering.__init__s"  !&)  r3c.|jjSrzr embed_tokens)r/s r2get_input_embeddingsz0MistralForQuestionAnswering.get_input_embeddings szz&&&r3c&||j_yrzr)r/values r2set_input_embeddingsz0MistralForQuestionAnswering.set_input_embeddings#s"' r3 input_idsrG position_idsrrstart_positions end_positionsrRoutput_hidden_statesrKc |j||||||| } | j} |j| } | jdd\}}|j dj }|j dj }d}|||j ||||fi| }t|||| j| jS)a start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. )rGrrrrRrrrM)rN)loss start_logits end_logitsrE attentions) rlast_hidden_state qa_outputssplitsqueezerc loss_functionrrEr)r/rrGrrrrrrRrrJoutputssequence_outputlogitsrrrs r2rmz#MistralForQuestionAnswering.forward&s0,0:: )%+'/!5,6, "331#)<<r<#: j#++B/::< ''+668   &=+D%4%%lJQ^ibhiD+%!!//))   r3) NNNNNNNNN)r4r5r6base_model_prefixr'rrrrorq FloatTensorrr rrrrmr7r8s@r2rrs '( 156:37KO596:48,0/33 E,,-3 !!2!233 u//0 3 "%tE4E4E/F(F"GH 3   1 12 3 "%"2"233   0 013 $D>3 'tn3  &3 r3r)3typingrrrrrrotorch.utils.checkpointr cache_utilsr r r modeling_attn_mask_utilsr modeling_flash_attention_utilsrmodeling_outputsrrmodeling_utilsrprocessing_utilsrutilsrllama.modeling_llamarrrrrrrrrrconfiguration_mistralr get_loggerr4r]_CHECKPOINT_FOR_DOCr"r:rsrxrrrrrr3r2rs99 AA>BU5&   1   H %1YY8)~8)v&+&]:]@ )  $?  'E A ";A r3