rh,dZddlZddlmZmZmZmZddlZddlZddlm Z ddl m Z m Z m Z ddlmZddlmZmZdd lmZdd lmZdd lmZdd lmZmZmZdd lmZmZddl m!Z!ddl"m#Z#m$Z$ddl%m&Z&m'Z'ddl(m)Z)e'rddl*m+Z+m,Z,ddl-m.Z.nd\Z.Z,Z+e&r ddl/m0Z0m1Z1nd\Z1Z0e2e.e,e0e1e+fZ3e$jhe5Z6Gdde jnZ8dejrde:dejrfdZ;GddeZ< d?d e jnd!ejrd"ejrd#ejrd$eejrd%e=d&e=fd'Z>Gd(d)e jnZ?Gd*d+e jnZ@Gd,d-e jnZAGd.d/e jnZBGd0d1e jnZCGd2d3e jnZDe#Gd4d5eZEe#Gd6d7eEZFGd8d9eEeZGe#d:;Gd<d=eEZHgd>ZIy)@zPyTorch Zamba model.N)AnyCallableOptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache DynamicCache)GenerationMixin)AttentionMaskConverter)FlashAttentionKwargs)BaseModelOutputWithPastCausalLMOutputWithPast SequenceClassifierOutputWithPast)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringlogging)is_causal_conv1d_availableis_mamba_ssm_available) ZambaConfig)mamba_inner_fnselective_scan_fn)selective_state_update)NNN)causal_conv1d_fncausal_conv1d_updateNNc,eZdZdfd ZdZdZxZS) ZambaRMSNormct|tjt j ||_||_y)z; ZambaRMSNorm is equivalent to T5LayerNorm N)super__init__r Parametertorchonesweightvariance_epsilon)self hidden_sizeeps __class__s {/var/www/html/ai-insurance-compliance-backend/venv/lib/python3.12/site-packages/transformers/models/zamba/modeling_zamba.pyr(zZambaRMSNorm.__init__@s1 ll5::k#:; #c"|j}|jtj}|j dj dd}|tj ||jzz}|j|j|zS)NT)keepdim) dtypetor*float32powmeanrsqrtr-r,)r. hidden_states input_dtypevariances r2forwardzZambaRMSNorm.forwardHsy#)) %((7  $$Q',,R,>% Ht?T?T4T(UU {{]--k:::r3c^t|jjd|jS)Nz, eps=)tupler,shaper-r.s r2 extra_reprzZambaRMSNorm.extra_reprOs* ))*+6$2G2G1HIIr3)gư>)__name__ __module__ __qualname__r(rArF __classcell__r1s@r2r%r%?s$;Jr3r%r>n_repreturnc|j\}}}}|dk(r|S|dddddddddfj|||||}|j|||z||S)z This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) rN)rDexpandreshape)r>rLbatchnum_key_value_headsslenhead_dims r2 repeat_kvrUTso 2?1D1D.E h z!!Qa"23::5BUW\^bdlmM  (;e(CT8 TTr3c eZdZdZdZdZdZejdfdZ dZ de de ejejffdZ dd ejd ejde d eeeefde ejejff d Zd ej(fdZddee de fdZde e eje ejffdZeddee e ej2ddfdZy)ZambaHybridDynamicCachea A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache (which has a constant shape regardless of seq_len). This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states` and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`, while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors). For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors), while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`, and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`. NFc ,||_d|_|j|_d|_|j|j z|_|j|_|j|_ |j|_ g|_ g|_ g|_i|_i|_i|_t%|j&D]}|xjt)j*||j |j||gz c_ ||j|j |jz|jf}|xjt)j*|||gz c_ |j|dk(s|jj-|t%|j&Dcgc]}t)j.gg|z|c}|_t%|j&Dcgc]}t)j.gg|z|c}|_ycc}wcc}w)NFdevicer8hybridrZ)r8is_compileablelayers_block_typehas_previous_state mamba_expandr/intermediate_size mamba_d_statessm_state_size mamba_d_convconv_kernel_size n_mamba_heads conv_states ssm_statestransformer_layers_modules _parameters_buffersrangenum_hidden_layersr*zerosappendtensor key_cache value_cache)r.config batch_sizer8rZi cache_shape_s r2r(z ZambaHybridDynamicCache.__init__rs #!'!9!9"'!'!4!4v7I7I!I$22 & 3 3#11"$  v//0 2A    J(>(>@U@U^dlqr!  ""&&$*<*<<## K OO Ke TU UO%%a(H4''..q1 2SXX^XpXpRqrQ%,,tj'8HrTYZ`ZrZrTstqELL" ):6Jtsts "H !"Hc,t|jSN)lenrrrEs r2__len__zZambaHybridDynamicCache.__len__s4>>""r3 layer_idxrMc>|j||j|fSrz)rrrsr.r}s r2 __getitem__z#ZambaHybridDynamicCache.__getitem__s!~~i($*:*:9*EEEr3 key_states value_states cache_kwargsc|j|jddk(r||j|<||j|<nftj|j||gd|j|<tj|j||gd|j|<|j||j|fS)Nr6rr5dim)rrrDrsr*cat)r.rrr}rs r2updatezZambaHybridDynamicCache.updates >>) $ * *2 .! 3(2DNN9 %*6D  Y '(- 4>>)3Lj2Y_`(aDNN9 %*/))T5E5Ei5PR^4_ef*gD  Y '~~i($*:*:9*EEEr3beam_idxctt|jD]S}|j|j}|j|j d|j ||j|<|j |j}|j |j d|j ||j |<|j|j}|j|j d|j ||j|<|j|j}|j|j d|j ||j|<Vy)zDReorders the cache for beam search, given the selected beam indices.rN) rmr{rrrZ index_selectr9rsrgrh)r.rr}rZs r2 reorder_cachez%ZambaHybridDynamicCache.reorder_cachesDs4>>23 iI^^I.55F(,y(A(N(NqRZR]R]^dRe(fDNN9 %%%i077F*.*:*:9*E*R*RSTV^VaVabhVi*jD  Y '%%i077F*.*:*:9*E*R*RSTV^VaVabhVi*jD  Y '__Y/66F)-)C)P)PQRT\T_T_`fTg)hDOOI & ir3c||jvr|jdn|}t|j|kry|j|jdS)zYReturns the sequence length of the cached states. A layer index can be optionally passed.r)rir{rrrDrs r2get_seq_lengthz&ZambaHybridDynamicCache.get_seq_lengthsR3<4CZCZ2ZD++A.`i t~~ ) +~~i(..r22r3ctdNz@ZambaHybridDynamicCache does not have a legacy cache equivalent.NotImplementedErrorrEs r2to_legacy_cachez'ZambaHybridDynamicCache.to_legacy_caches!"deer3past_key_valuesrctdrr)clsrs r2from_legacy_cachez)ZambaHybridDynamicCache.from_legacy_caches!"deer3rz)r)rGrHrI__doc__rrrsr]r*float16r(r|intrCTensorrrdictstrrr LongTensorrrr classmethod FloatTensorrr3r2rWrW`sM IKN16tu@#FSFU5<<3M-NF26 FLLFllF F tCH~. F u||U\\) * F$ ie&6&6 i3 3c3fuU\\':E%,,$>&B\B\$\!'-'E'E$ )d2 !'!9!9ii < X>X[_[h[h>hotu ii < X>X[_[h[h>hotu ii < X>X[_[h[h>hotu ii : :T]] JFL^L^ejk r3r>rpast_key_valuerrMc |jdd}g|d|j}|j|j|j dd}|j |j|j dd} |j |j|j dd} ||j| | |\} } t} |jjdk7rt|jj} | ||| | |f|jsdn |j|jd|\} } | jg|dj!} |j#| } | | fS)Nr6rr5eager)rr)rDrTrviewrrrrrrt_attn_implementationrrrrrPrr)r.r>r}rrr input_shape hidden_shape query_statesrrattention_interfacerrs r2rAzZambaAttention.forwardsu$))#2. 88b8$--8 {{=166|DNNqRST [[/44\BLLQPQR {{=166|DNNqRST  %'5'<'Q5RR S#)r3rcleZdZdZdeffd Z d dejdefdZ d defdZ d defdZ xZ S) ZambaMambaMixeruE Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, and is why Mamba is called **selective** state spaces) This module differs from `transformers.models.mamba.modeling_mamba.MambaMixer` in two ways: - Added multi-head: the output of `self.in_proj` is split into `self.n_mamba_heads` heads, and each head undergoes an independent forward pass, identical to the original `MambaMixer`, up until the pre-activations of `self.out_proj`. The pre-activations, coming from different mamba heads, are then concatenated and fed into `self.out_proj`. rtc t|||_||_|j|_|j |_|j|_|j|jz|_ |j|_ |j|_ |j|jz|_|j|_|j"|_t'j(|j|j|j |j|j|jdz |_|j,|_t0|j,|_|j4|_t'j8|j|jdz|j$|_t'j<t?j@|j|j|j dzz|j|_!t'j<t?j@|j|j|jdz dz|jdzz |_"t'j<t?j@|j|j|_#t?jHd|j dzt>jJdddf}|jM|jdjO}t'j<t?jP|jS|j|jd|_*t'j<t?jV|j|j|_,t'j8|j|j|j$|_-t\st^jadyy) Nr) in_channels out_channelsr kernel_sizegroupspaddingr5rg?r8r6apThe fast path is not available because on of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)` is None. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d. If you want to use the naive implementation, set `use_mamba_kernels=False` in the model config)1r'r(rtr}r/rbrcrdrer`ra mamba_dt_ranktime_step_rankrfmamba_head_dimmamba_conv_bias use_conv_biasmamba_proj_biasuse_biasrConv1dconv1dhidden_mamba_act activationr actuse_mamba_kernelsuse_fast_kernelsrin_projr)r*ro x_proj_weightdt_proj_weight dt_proj_biasaranger:rOrlogrPA_logr+Dout_projis_fast_path_availablelogger warning_once)r.rtr}Ar1s r2r(zZambaMambaMixer.__init__6s   "!--$22 & 3 3!'!4!4v7I7I!I$22#11"448J8JJ#33.. ii..//##--))))A-   !11&112 & 8 8yy!1!143I3IA3MTXTaTab  \\ KK""##d&9&9A&==##  !ll [[++T-@-@$BUBU VY\ \ !!3& '  LLT5G5GI\I\)]^ LLD//!35== I$PQ' R HHT++R 0 ; ; =\\%))A,"6"6t7I7I4K^K^`b"cd ejj););T=P=PQR $"8"8$:J:JQUQ^Q^_ %   ^ &r3r> cache_paramsc |j\}}}|duxr|jxr|dk(}|j|jdd}|j |dd|j dd\}} |j dj}| j d} | j||jd|jdd} |jjj |jjjd|jjjd} |ret|j d|j|j| |jj |j"}|j%d}n|,t'j(|dk(s||j%dz}|dt*j,j/||j0|jdz df} |j|jj3| t5|| |jj |j"}|,t'j(|dk(s||j%dz}|jd|j|j6|jdd}|j8dddddddf|zjdd} t'j:| |j<|j>|j>gd\} }}|j@dddf| jddz}t'jB|jDjG }|jH|jHjGnd}t'jJ|d|f|jL|jN}|rtQ|jD]}tS|jT|jdd|f||d df||d df||||dddf||dddf|jV|| |d df||d j%d}t'jX||fd}nAt'jJ|d|j6|j>f|jL|jN}tQ|jD]}t[||||||||jdd||jdd|jV|jG| |||d d \}}t'jX||fdj}t'jX||j%dfd}|*|(|jT|jj3||j]|jdd}|S) Nrr5r6rr)rrrY.T) dt_softplus)delta_softplusreturn_last_state)/rDr_rrrchunksqueezerrPrfrr,sizer"rgr}rr unsqueezer*allrrpadrecopy_r!rrsplitrrcrexprfloatremptyrZr8rmr rhrrrr)r.r>rrruseq_lenrxuse_precomputed_statesprojected_statesgate conv_weightsrgssm_parameters time_stepBCdiscrete_time_steprtime_proj_bias scan_outputsn scan_outputs_ ssm_state ssm_state_contextualized_statess r2cuda_kernels_forwardz$ZambaMambaMixer.cuda_kernels_forwardss"/!4!4 GQ!-T!9!nl>]>]!nbimnbn << 6@@AF.33JAwOUUVW]^U_ t%--a0;;= ||A||J(:(:BHRRSTVWX{{))..t{{/A/A/F/Fq/I4;;K]K]KbKbcdKef !0%%b)((8    M*33B7M)%))Na>{K,]L$++JZJZgkgvgvwM)%))NajA| |jB| jdz df} | |j&|j<|j7|j/| dd|f} | t)j:|dk(s| |dd| jd dfj9dz} nt)jD||j|jF|jHf| j"|} |,t)j:|dk(s| |j9dz} |j7|j/| dd|f} |,t)j:|dk(s| |j9dz} | jd|j|jF|jdd} |jJdddddddf| zjdd }t)jL||jN|jH|jHgd\}}}|jPdddf|jdd z|jRddddddfz}t<j>jU|}t)jV|jXj[ }t)jV|ddddddddf|dddddddddfz}|dddddddddf|dddddddddfj[z}|| dddddddddfj[z}g}t]|D]}|dddddd|ddfjdd| z|dddddd|ddfjddz} t)j^| jddj!||dddd|ddfj9d}|ja|dddddddft)jb|d}|| |jdddddddfzz}||j7| z}| r| |j|j<|jg|jddj|d|jdd}|S) Nrr5r6rr)shiftsdims.rYr)4rDr8rrrrrrrPrf isinstancerWrhr}rcloner9rZr_rgr*rollsumrr,rrrrrrrrrerorrcrrrrrsoftplusrrrrmrrpstackrr)r. input_statesrrrurrxr8r r>r  use_cacher conv_stater rrrrr discrete_A discrete_BdeltaB_urrv scan_outputrs r2 slow_forwardzZambaMambaMixer.slow_forwards!-!3!3 GQ""<< 5??1E.33JAwOUUVW]^U_ t%--a0;;= ||A||J(:(:BHRRSTVWX|-DE 00@FFqIZW}}(33DNNCIIK (33DNNC ! ]%9%9:I//qL ,,T^^<BB1ES)55dnnE "ZZ 2BG '4Q1W'= 1a8$;E ((8 % *t{{7I7I!QPQ'7R*RXZ [ %%!T[[%5%55M $ 7 : :5 A K KB O !-eiiRS@S6T$1N1}GZGZ[]G^F^F`C`4a4k4klm4n$nM]]..}t?T?TWdWjWjkmWn?npq>rs ;E ((8 $])CC'M)R S !-eiiRS@S6T$1N1}GZGZ[]G^F^F`C`4a4k4klm4n$nM T//1D1DdFYFYZ$++I )%))NaN ? E E GG w 9A"1aAq=1;;AqAIMPXYZ\]_`bcefYfPgPqPqrsuvPwwI,,y':':1a'@'C'CE'JAaQRTUWXjMLcLcdfLghK    Aq!QJ 7 8 9kk,B7 !]TVVAtQ=1.2.0' and 'pip install mamba-ssm', or set use_mamba_kernels=False in the model's config.)r)rrrrZtype ValueErrorrr*)r.r>rrs r2rAzZambaMambaMixer.forward1sm  )V4;M;M;T;T;Y;Y-Y i ,,]LYg,h h   ^ \\r3r#) rGrHrIrrr(r*rrWrr*rArJrKs@r2rr)sX ;{;|im_%"\\_%9P_%B[%7N[%z ]3J ]r3rc$eZdZfdZdZxZS)ZambaMLPct|||_|j|_|j|_t j |j|jd|_t j |j|jd|_t j |j|jd|_ t|j|_ yNFr) r'r(rtr/rarr gate_projup_proj down_projr hidden_actact_fnr.rtr1s r2r(zZambaMLP.__init__?s  !--!'!9!94#3#3T5K5KRWXyy!1!143I3IPUV 4#9#94;K;KRWXV../ r3c|j|j|j||j|z}|Srz)r5r7r3r4)r.xr5s r2rAzZambaMLP.forwardIs6NN4;;t~~a/@#ADLLQRO#ST r3)rGrHrIr(rArJrKs@r2r0r0>s 0r3r0ceZdZd dedeeffd Z ddejdejdedeejdee dee d ee d e e d e ejee ejejfffd ZxZS)ZambaAttentionDecoderLayerrtr}ct|t|||_t ||_t |j|j|_ t |j|j|_ y)Nr0) r'r(r self_attnr0 feed_forwardr%r rms_norm_epsinput_layernormr/pre_ff_layernormrs r2r(z#ZambaAttentionDecoderLayer.__init__Os_ ' :$V,+F,H,HfNaNab ,V-?-?VEXEX Yr3r>original_hidden_statesrroutput_attentionsr$rrMc tj||gd}|j|}|jd||||||d|\}} |j |}|j |}|f} |r| | fz } | S)a Args: hidden_states (`torch.FloatTensor`): output of previous Mamba layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output of shape `(batch, seq_len, embed_dim)`. This is concatenated with `hidden_states` (which is the output of the previous (mamba) layer). The concatenated tensor is then used as input of the pre-attention RMSNorm (see fig. 2 in https://huggingface.co/papers/2405.16712). layer_idx (`int`): layer_idx in the forward pass. Used to distinguish Zamba's tied transformer layers. attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`ZambaHybridDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. r6r)r>r}rrrEr$r)r* concatenaterBr?rCr@) r.r>rDr}rrrEr$rself_attn_weightsoutputss r2rAz"ZambaAttentionDecoderLayer.forwardWs>))=:P*QWYZ ,,]; +94>>, '))/ , , ( (--m< ))-8  "  )+ +Gr3rz)NNFF)rGrHrIrrrr(r*rrWboolrrrCrrArJrKs@r2r<r<NsZ{Zx}Z26<@,1$)3||3!& 3 3 !. 3 !!89 3$D>3D>3-.3 u  (51B1BEDUDU1U+V"WW X3r3r<cteZdZdedeffd Z ddejdeejdeedeejdeejdee d ee d ee d eejd eejd e ejee ejejfffdZxZS)ZambaMambaDecoderLayerrtr}ct|t|||_t |j |j |_||_y)N)rtr}r>) r'r(rmambar%r/rArBr}rs r2r(zZambaMambaDecoderLayer.__init__s> $FiH +F,>,>FDWDWX"r3r>rDrrrrEr$cache_positiontransformer_hidden_statesrMc  |} | || zn|}|j|}|j|||}d} | |z}|f}|r|| fz }|r||fz }|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`ZambaHybridDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. N)r>rr)rBrN)r.r>rDr}rrrrEr$rOrPrresidualrHrIs r2rAzZambaMambaDecoderLayer.forwards<! :S9^M5 5dq ,,];  '')# !!=0  "  )+ +G  ( (Gr3) NNNNNFFNN)rGrHrIrrr(r*rrrWrJrrCrrArJrKs@r2rLrLs#{#s#:>#'15.2<@,1$)59<@:||:!) 6:C= : !. : ell+ :!!89:$D>:D>:!!1!12:$,ELL#9: u  (51B1BEDUDU1U+V"WW X:r3rLcleZdZdedej deffd Z ddejde ejde e de ejd e ejd e e d e e d e e d e ejdeej e eej ej fffdZxZS)ZambaHybridLayer shared_transflinearrNcLt|||_||_||_yrz)r'r(rUrV mamba_decoder)r.rUrVrNr1s r2r(zZambaHybridLayer.__init__s% * "r3r>rDr}rrrrEr$rOrMc |j|||||||| } | d} |r| d} |j| } |j|| ||||| } |r | d f| ddz} | S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` original_hidden_states (`torch.FloatTensor`): word embedding output that will be concatenated with hidden activations to form the input of the shared transformer layer. layer_idx (`int`): layer number. attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, sequence_length)` where padding elements are indicated by 0. past_key_value (`ZambaHybridDynamicCache`, *optional*): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. )rDr}rrrEr$rOrr)rPrrrEr$rOr5N)rUrVrX) r.r>rDr}rrrrEr$rO layer_outputsrPrHs r2rAzZambaHybridLayer.forwards>** #9&)/)+  %2!$4!  -a 0 $(KK0I$J!** &?))/)+  *1-/@AMRSRTDUUMr3)NNNNNFFN)rGrHrIr<rrrLr(r*rrrrWrJrrCrrArJrKs@r2rTrTs #&@#"))#\r#:>#'15.2<@,1$)59>||>!) 6>C= > !. > ell+ >!!89>$D>>D>>!!1!12> u  (51B1BEDUDU1U+V"WW X>r3rTc>eZdZUeed<dZdZddgZdZdZ dZ dZ dZ y ) ZambaPreTrainedModelrtmodelTr<rLrFcT|jj}t|tjtj frY|j jjd||j%|jjjyyt|tjrf|j jjd||j2|j j|jjyyt|tr&|j jjdyt|tr|j jjd||jj"dz}tj$j'|j(| ||jj*|jj,z|jj.z}t1j2t1j4|jj.|t7j8|jj:t7j8|jj<z zt7j8|jj<zj?|jj@}|t1j8t1jB|  z}|jDjjG|t1jHd|jJdzt0jLdddf}|jO|jPdjS}|jTjjGt1j8|jW|j.|jXd|jZjjdyy) Nr)r<stdg?r)minrrr6).rtinitializer_rangerrrrr,datanormal_rzero_ Embedding padding_idxr%fill_rrrinituniform_rr`r/rfr*rrandmathr time_step_max time_step_minclamptime_step_floorexpm1rrrrcr:rOrarrrPrr)r.rr_ dt_init_stdrdtinv_dtrs r2 _init_weightsz"ZambaPreTrainedModel._init_weights%skk++ fryy"))4 5 MM   & &CS & 9{{&   &&('  - MM   & &CS & 9!!- ""6#5#56<<>.  - MM   $ $S )  0  % % - -3C - @++33T9K GG  V22[L+ N![[55 8O8OOSWS^S^SlSllN 4;;44nE88DKK556$++B[B[9\\^((4;;4456e 33e4  %))U[["%5$566F    $ $ * *6 2 Q 5 5 9OPTVWPWXA1126AACA LL   # #EIIaL$8$89M9MvOdOdfh$i j HHMM   $%1r3N) rGrHrIr__annotations__base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa _is_statefulrtrr3r2r\r\s; &*#57OP"3 NL%r3r\c,eZdZdZdeffd Ze ddeejdeejdeejdee deejd ee d ee d ee d ee d eejdeeeffdZdZxZS) ZambaModelz Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`ZambaDecoderLayer`] Args: config: ZambaConfig rtc t |||j|_|j|_t j |j|j|j|_t|}g}g}|j|_ t|jD]}|j|dk(r|jt||2|j|dk(sE|jt j|j j|j jd|jt||t#|}t#|}g}g|_t'|jD]\}}|dk(r_d|d} gd} g|j$| D cgc]} | | z c} |_|jt)|t+|t+|j|jt+|t j,||_|j0|_t3|j|j4 |_d|_|j;ycc} w) NrN)r}r[Frzlayers..) z%shared_transf.self_attn.q_proj.weightz%shared_transf.self_attn.k_proj.weightz%shared_transf.self_attn.v_proj.weightz%shared_transf.self_attn.o_proj.weightz+shared_transf.feed_forward.gate_proj.weightz)shared_transf.feed_forward.up_proj.weightz+shared_transf.feed_forward.down_proj.weightz$shared_transf.input_layernorm.weightz%shared_transf.pre_ff_layernorm.weightr>)r'r( pad_token_idrf vocab_sizerrer/ embed_tokensr<r^rmrnrprLrrtiter_tied_weights_keys enumeraterTnext ModuleListlayersrr%rAfinal_layernormgradient_checkpointing post_init) r.rtblock mamba_layers linear_layersrvrlayer_id layer_type prefix_name tied_keysrr1s r2r(zZambaModel.__init__Os&  !.. ++LL):):F,>FDWDWX&+# Gos. I7 input_idsr position_idsr inputs_embedsr$rEoutput_hidden_states return_dictrOrMc d||n|jj}||n|jj}||n|jj}| | n|jj} |du|duz r t d|j r%|jr|rtjdd}||j|}|} tj|} |r|tjd| .tj| jd| j} || j!d}|j#||| } |rdnd}|rdnd}t%|j&D]r\}}|r|| fz }|j r1|jr%|j)|j*| | ||| |||| }n|| | ||| |||| }|d} |sd|dj||dfz }t|j-| } |r|| fz }|r|j.sd |_t1| |r|nd|| }| r|S|j3S) NzaYou cannot specify both input_ids and inputs_embeds at the same time, and must specify either onezX`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.Fz{Zamba requires an initialized `ZambaHybridDynamicCache` to return a cache. None was provided, so no cache will be returned.rr\rr)rDr}rrrrEr$rOT)last_hidden_staterr> attentions)rtrErr$use_return_dictr.rrrrrr*rrrDrZr_update_causal_maskrr_gradient_checkpointing_func__call__rr_rto_tuple)r.rrrrrr$rErrrOr>rDrall_hidden_statesall_self_attnsr}layerrZoutputs r2rAzZambaModel.forward~s2C1N-TXT_T_TqTq$8$D $++JjJj "+!6IDKKk<$D>k<'tnk<d^k<!!1!12k< u-- .k<k<\!r3r~ceZdZdeffd ZdZdZe ddee jdee jdee jdee d ee jd ee jd eed eed eedeedee jdeee jfdeeeffdZ ddZxZS)ZambaForCausalLMrtc$t||t||_dg|jj|_|j |_t j|j|j d|_ |jy)Nzlm_head.weightFr) r'r(r~r]rrrrr/lm_headrr8s r2r(zZambaForCausalLM.__init__so  ' #3"Tdjj6S6S"T ++yy!3!3V5F5FUS  r3c||_yrzr])r.decoders r2 set_decoderzZambaForCausalLM.set_decoders  r3c|jSrzrrEs r2 get_decoderzZambaForCausalLM.get_decoder s zzr3rrrrrlabelsr$rErrrOlogits_to_keeprMc ||n|jj}| | n|jj} | | n|jj} |j |||||||| | |  }|d}t | t r t| dn| }|j|dd|ddf}d}||j|||jfi| }| s|f|ddz}||f|zS|St|||j|j|jS)ah labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (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 >>> from transformers import AutoTokenizer, ZambaForCausalLM >>> model = ZambaForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1") >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba-7B-v1") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```N) rrrrrr$rErrOrrrlosslogitsrr>r)rtrErrr]rrslicer loss_functionrrrr>r)r.rrrrrrr$rErrrOrrrIr> slice_indicesrrrs r2rAzZambaForCausalLM.forward#sLP2C1N-TXT_T_TqTq%9$D $++JjJj &1%SV8W~ot4]k mA}a,?@A  %4%%ffdooPPDY,F'+'7D7V# CV C%#33!//))   r3c t|du} | sZ||d|jdk\r|dd|jd df}nc|jd|jdk7rD|dd|f}n:t|j|jd|j|j}|T|R|j j ddz }|j|dk(d| s|dd|jd df}|| rd|i} nd|ji} | j|||||jj|d| S)Nr6rr)r8rZrr)rrr$rrrO) rDrWrtr8rZlongcumsum masked_fill_rrnum_logits_to_keep) r.rrrrrOrr$r empty_past_kv model_inputss r2prepare_inputs_for_generationz.ZambaForCausalLM.prepare_inputs_for_generationusc(4/ )!"%);;%a.*>*>q*A)A)C&CD #~';';A'>>%a&78 5 Y__Q/tzz$++O  %,*>)..077;a?L  % %n&91 = +A 0B/B/D,DE   $+];L')=)=)?@L ,#2&"0"&++"@"@"0   r3) NNNNNNNNNNNr)NNNNNT)rGrHrIrr(rrrrr*rrrWrrJrrrCrrArrJrKs@r2rrsk{151537=A59-1$(,0/3&*5934O E,,-O !.O u//0 O ""9: O   1 12 O ))*O D>O $D>O 'tnO d^O !!1!12O c5<</0O  u,, -O O h9r3ra The Zamba Model with a sequence classification head on top (linear layer). [`ZambaForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-2) 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). ) custom_introc@eZdZfdZe ddeejdeejdeejdee e e ejfdeejdeejdee d ee d ee d ee d e eeffd ZxZS)ZambaForSequenceClassificationct|||j|_t||_|jj |_t j|j|jd|_ |jyr2) r'r( num_labelsr~r]rrrr/scorerr8s r2r(z'ZambaForSequenceClassification.__init__se   ++' "&**"?"?YYv114??O  r3rrrrrrr$rErrrMc | | n|jj} |j|||||||| |  } | d} |j| } ||jd}n|jd}|jj |dk7r t d|jj d}n|||jj k7j| jtj}tj|jd| jtj}||zjd}n.d}tj|jj d| tj|| j |f}d}||j| j}|jj"|j$dk(rd |j_nl|j$dkDrL|j&tj(k(s|j&tj*k(rd |j_nd |j_|jj"d k(rIt-}|j$dk(r&||j/|j/}n|||}n|jj"d k(r=t1}||j3d|j$|j3d}n,|jj"d k(rt5}|||}| s|f| ddz}||f|zS|St7||| j8| 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). N)rrrrr$rErrrrz=Cannot handle batch sizes > 1 if no padding token is defined.r6rYz will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`r\ regressionsingle_label_classificationmulti_label_classificationr)rtrr]rrDrr.r9rZr*int32rargmaxrrr1rG problem_typerr8rrr rr rrrrr>r)r.rrrrrrr$rErrtransformer_outputsr>rrulast_non_pad_token non_pad_mask token_indices pooled_logitsrloss_fctrs r2rAz&ZambaForSequenceClassification.forwards(&1%!F!Fr!J !#    >>**+,ZZ  u||Jv}}MOaab   YYv}}-F{{''/??a'/;DKK,__q(fllejj.HFLL\a\e\eLe/LDKK,/KDKK,{{''<7"9??a'#M$9$9$;V^^=MND#M6:D))-JJ+- 2 22t GUWY))-II,. v6#%(;AB(??F)-)9TGf$ Ev E/ /??-;;*55   r3r)rGrHrIr(rrr*rrrr listrrJrCrrArJrKs@r2rrs151537KO59-1$(,0/3&*[ E,,-[ !.[ u//0 [ "%tE4E4E/F(F"GH [   1 12 [ ))*[ D>[ $D>[ 'tn[ d^[  u66 7[ [ r3r)rrr~r\)r)Jrrktypingrrrrr*torch.utils.checkpointrtorch.nnrr r activationsr cache_utilsr r generationrmodeling_attn_mask_utilsrmodeling_flash_attention_utilsrmodeling_outputsrrrmodeling_utilsrrprocessing_utilsrutilsrrutils.import_utilsrrconfiguration_zambar&mamba_ssm.ops.selective_scan_interfacerr+mamba_ssm.ops.triton.selective_state_updater causal_conv1dr!r"rr get_loggerrGrModuler%rrrUrWrrrrr0r<rLrTr\r~rr__all__rr3r2rsH( 11 AA!.)>BqqF&,T,XR@P=-~DD-7**.0@BVXfg   H %J299J* UU\\ U# U%,, Uffeff`% II% <<% % << % U\\* %  %%4C)RYYC)LQ]biiQ]j ryy  <<~ARYYAHEryyEP)%?)%)%XG%GGV\+_\~  g %9g  g T gr3