a hL @sddlZddlZddlmZddlmZmZmZmZddl Z ddl m Z ddl m Z m Z mZddlmZddlmZdd lmZdd lmZdd lmZdd lmZmZmZdd lmZmZddl m!Z!m"Z"ddl#m$Z$ddl%m&Z&m'Z'ddl(m)Z)ddl*m+Z+m,Z,ddl-m.Z.e,r6ddl/m0Z0ddl1m2Z2m3Z3n d\Z0Z2Z3e+rZddl4m5Z5m6Z6nd\Z6Z5e'7e8Z9Gddde j j:Z;Gddde j:Ze j?e@e j?d"d#d$ZAdKe j:e j?e j?e j?ee j?eBeBd&d'd(ZCd)d*ZDdLd+d,ZEGd-d.d.e j:ZFe j?e@d/d0d1ZGd2d3ZHd4d5ZIeJe0e5e6fZKGd6d7d7e j:ZLGd8d9d9e j:ZMGd:d;d;e j:ZNGdd?d?e j:ZPGd@dAdAe"ZQe&GdBdCdCeQZRGdDdEdEeQeZSe&dFdGGdHdIdIeQZTgdJZUdS)MN)cycle)AnyCallableOptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache)GenerationMixin)AttentionMaskConverter)FlashAttentionKwargs)BaseModelOutputWithPastCausalLMOutputWithPast SequenceClassifierOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringlogging)deprecate_kwarg)is_causal_conv1d_availableis_mamba_ssm_available) Zamba2Config)selective_state_update)mamba_chunk_scan_combined mamba_split_conv1d_scan_combined)NNN)causal_conv1d_fncausal_conv1d_update)NNcs(eZdZdfdd ZdddZZS) Zamba2RMSNormGatedư>cs,ttt||_||_||_dSN) super__init__r Parametertorchonesweightvariance_epsilon group_size)self hidden_sizer/eps __class__f/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/zamba2/modeling_zamba2.pyr)=s zZamba2RMSNormGated.__init__Nc Cs|j}|tj}|dur2|tj|tj}|j^}}||j}|j g|||jR}| dj ddd}|t ||j }|j g|||jR}|j||SNT)Zkeepdim)dtypetor+float32r functionalsilushaper/viewpowmeanrsqrtr.r-) r0 hidden_statesgate input_dtypeZ prefix_dimsZlast_dimZ group_countZhidden_states_groupvariancer5r5r6forwardCs   zZamba2RMSNormGated.forward)r&)N)__name__ __module__ __qualname__r)rH __classcell__r5r5r3r6r%<sr%cs.eZdZdfdd ZddZddZZS) Zamba2RMSNormr&cs&ttt||_||_dS)z< Zamba2RMSNorm is equivalent to T5LayerNorm N)r(r)rr*r+r,r-r.)r0r1r2r3r5r6r)Rs zZamba2RMSNorm.__init__cCsJ|j}|tj}|djddd}|t||j}|j||Sr7) r:r;r+r<rArBrCr.r-)r0rDrFrGr5r5r6rHZs  zZamba2RMSNorm.forwardcCst|jjd|jS)Nz, eps=)tupler-r?r.r0r5r5r6 extra_repraszZamba2RMSNorm.extra_repr)r&)rIrJrKr)rHrPrLr5r5r3r6rMQsrMc@seZdZdZdZejdfeeej e e dddZ ddZ eeejejfd d d Zdejejee ee efeejejfd d dZejdddZde eed ddZeejejejdddZddZdS)Zamba2HybridDynamicCachea 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)`. FN)config batch_sizer:devicec s||_|j|_d|_t|j|j|_|j|_|j |_ |j |_ g|_ i|_ i|_i|_i|_i|_t|jD]n}tj|jd|j|j|j |d|j|<tj|j |j|j|d|j|<|j|dkrl|j |qlfddt|jD|_fddt|jD|_dS)NFr8rTr:hybridcs g|]}tjggdqSrTr+Ztensor.0_rSrTr5r6 z5Zamba2HybridDynamicCache.__init__..cs g|]}tjggdqSrWrYrZr]r5r6r^r_)r:layers_block_typehas_previous_stateint mamba_expandr1intermediate_size mamba_d_statessm_state_size mamba_d_convconv_kernel_size n_mamba_headstransformer_layersZ_modules _parameters_buffers conv_states ssm_statesrangenum_hidden_layersr+zeros mamba_ngroups mamba_headdimappend key_cache value_cache)r0rRrSr:rTir5r]r6r)us8  z!Zamba2HybridDynamicCache.__init__cCs t|jSr')lenrurOr5r5r6__len__sz Zamba2HybridDynamicCache.__len__) layer_idxreturncCs|j||j|fSr')rurvr0rzr5r5r6 __getitem__sz$Zamba2HybridDynamicCache.__getitem__) key_states value_statesrz cache_kwargsr{cCsz|j|jddkr*||j|<||j|<n