U 9%eJ @sdZddlZddlmZddlmZddlmZmZm Z m Z ddl Z ddl Z ddl m Z ddlmZdd lmZdd lmZmZmZmZmZmZmZmZd d lmZeeZd Z dZ!d dddddddddg Z"da#ddZ$Gddde j%j&Z'd5ddZ(d6dd Z)Gd!d"d"e j*Z+Gd#d$d$e j*Z,Gd%d&d&e j*Z-Gd'd(d(eZ.eGd)d*d*eZ/eGd+d,d,eZ0d-Z1d.Z2ed/e1Gd0d1d1e.Z3ed2e1Gd3d4d4e.Z4dS)7zPyTorch RWKV model.N) dataclass)Path)ListOptionalTupleUnion)nn)CrossEntropyLoss)PreTrainedModel) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardis_bitsandbytes_availableis_ninja_availableis_torch_cuda_availablelogging) RwkvConfigzRWKV/rwkv-4-169m-pilerzRWKV/rwkv-4-430m-pilezRWKV/rwkv-4-1b5-pilezRWKV/rwkv-4-3b-pilezRWKV/rwkv-4-7b-pilezRWKV/rwkv-4-14b-pilezRWKV/rwkv-raven-1b5zRWKV/rwkv-raven-3bzRWKV/rwkv-raven-7bzRWKV/rwkv-raven-14bcsddlm}ttjjjddfdddD}tdk rNtj|krNdSt d|d d d d d ddd|g}|d||t t j k|da|t_dS)Nr)loadZkernelsrwkvcsg|] }|qSr).0fZ kernel_folderre/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/models/rwkv/modeling_rwkv.py Gsz(load_wkv_cuda_kernel..)z wkv_op.cppz wkv_cuda.cuzwkv_cuda_bf16.cuz2Loading CUDA kernel for RWKV at context length of .z -res-usagez--maxrregcount 60z--use_fast_mathz-O3z -Xptxas -O3z--extra-device-vectorizationz-DTmax=Zwkv_)namesourcesverboseZextra_cuda_cflags) Ztorch.utils.cpp_extensionrr__file__resolveparentrwkv_cuda_kernelmax_seq_lengthloggerinforZ get_verbosityDEBUG)context_lengthZ load_kernelZcuda_kernel_filesflagsrrrload_wkv_cuda_kernelAs*   r,c@s(eZdZedddZedddZdS) RwkvLinearAttentionNFc Cs|\}}} |tjkr0td|dtjd|| t| ddkrhtd|d| dt| dd |j|_|jjd ks|jjd ks|jjd ks|jjd krtd t |  }|jt j kr| }| }| }| }| }| }t j|t jd } |s|dk r|dkr^t j|| d t j|jt jd}|dddddfd8<nt jdd|Ddd }|jt jkrtj} ntj} | ||||| |n*|jt jkrtjntj} | ||||| |||||| |dk r ddt j|d ddD}| |j|fS)NzCannot process a batch with z+ tokens at the same time, use a maximum of z with this model. rzThe product of batch size (z) and hidden size (z") needs to be a round multiple of rcudazUCalling the CUDA kernel for wkv attention requires all tensors to be on CUDA devices. memory_formatr )dtypedevicer1籡*GcSsg|]}|dqSr4Z unsqueezersrrrrsz/RwkvLinearAttention.forward..)dimcSsg|]}|dqSr6)Zsqueezer8rrrrs)sizer%r& ValueErrorminr2 input_dtyper3typetorchexpfloat contiguousfloat16 empty_likecontiguous_formatzerosfloat32catbfloat16Zforward_with_state_bf16Zforward_with_stateZ forward_bf16forwardZsave_for_backwardchunkto) ctx time_decay time_firstkeyvaluestate return_stateZ batch_sizeZseq_len hidden_sizeoutputZ forward_funcrrrrKbsd         zRwkvLinearAttention.forwardc Cs|j}|j\}}}}}tj|tj|tjkr0tjntjd} tj|tjd} tj|tjd} tj|tjd} |tjkr||}|tjkrt j nt j } | |||||| | | | | | || || || |ddfS)N)r1r2r0)r>Z saved_tensorsr@rErFrJrHrDrBr%Z backward_bf16backwardrCrM)rNZg_outputZg_stater>rOrPrQrRrVZ g_time_decayZ g_time_firstZg_keyZg_valueZ backward_funcrrrrWs@ zRwkvLinearAttention.backward)NF)N)__name__ __module__ __qualname__ staticmethodrKrWrrrrr-as >r-FcCs|\}}}t|}|dkrztj|dddftjd} tj|dddftjd} tj|dddftjdd} n |\} } } t| }t|D]} |dd| f} |dd| f}t| | |}t| |}t| ||}|| ||}|| |}|||j |dd| f<t| || }t| ||}t| |}|| ||} || |} |} q|s|dk r| | | g}||fS)Nr)r2r5) r;r@Z zeros_likerHrArangerBmaximumrMr2)rOrPrQrRrSrT_Z seq_lengthrVZ num_stateZ den_stateZ max_stateZ current_index current_key current_valueZmax_for_outpute1e2 numerator denominatorZ max_for_staterrrrwkv_linear_attention_cpus4 "      recCsdtdd||||fD}|ddk}tdks8|s8|rLt||||||dSt||||||SdS)Ncss|]}|jjdkVqdS)r/N)r3r?)rtrrr sz(rwkv_linear_attention..rrSrT)anyr;r%rer-apply)rOrPrQrRrSrTZno_cudaZ one_tokenrrrrwkv_linear_attentions rkcs2eZdZd fdd Zd ddZd dd ZZS) RwkvSelfAttentionrcsFt||_tdk o"tj|jk}trbtrb|sbzt|jWnt k r`t dYnX||_ |j }|jdk r~|jn|}||_tt||_tt||_ttdd||_ttdd||_ttdd||_td|_tj||dd|_tj||dd|_tj||dd|_tj||dd|_dS)Nz9Could not load the custom CUDA kernel for RWKV attention.rrrrFZbias)super__init__configr%r&r*rrr, Exceptionr'r(layer_idrUattention_hidden_sizer Parameterr@emptyrOrP time_mix_keytime_mix_valuetime_mix_receptance ZeroPad2d time_shiftLinearrQrR receptancerV)selfrrrtZ kernel_loadedrUru __class__rrrqs.  zRwkvSelfAttention.__init__NcCs|ddkr4|dk r4|ddddd|jf}n:||}|dk rn|ddddd|jf|dddf<||j|d|j}||j|d|j}||j|d|j}||}||}t | |}|dk r |dddf|ddddd|jf<||||fSNrrrn) r;rtr|rxryrzrQrRr@sigmoidr~)rhiddenrSshiftedrQrRr~rrrextract_key_value!s (   (z#RwkvSelfAttention.extract_key_valueFc sj||d\}}}}|dk r.r4rhrrr )rtuplerkrOrPrtrV) rrrS use_cacher~rQrRZ layer_staterrrrrK4s*    zRwkvSelfAttention.forward)r)N)NF)rXrYrZrqrrK __classcell__rrrrrls rlcs(eZdZdfdd ZdddZZS) RwkvFeedForwardrcst||_||_|j}|jdk r,|jnd|j}td|_t t dd||_ t t dd||_ tj||dd|_tj||dd|_tj||dd|_dS)NrrmrFro)rprqrrrtrUintermediate_sizerr{r|rvr@rwrxrzr}rQr~rR)rrrrtrUrrrrrqIs  zRwkvFeedForward.__init__NcCs |ddkr4|dk r4|ddddd|jf}n:||}|dk rn|ddddd|jf|dddf<||j|d|j}||j|d|j}tt||}| |}t | |}|dk r|dddf|ddddd|jf<|||fSr) r;rtr|rxrzr@ZsquareZrelurQrRrr~)rrrSrrQr~rRrrrrKZs ( (zRwkvFeedForward.forward)r)NrXrYrZrqrKrrrrrrHsrcs&eZdZfddZdddZZS) RwkvBlockcsvt||_||_|dkr2tj|j|jd|_tj|j|jd|_ tj|j|jd|_ t |||_ t |||_dS)Nr)Zeps)rprqrrrtr LayerNormrUZlayer_norm_epsilonpre_lnln1ln2rl attentionr feed_forward)rrrrtrrrrqos  zRwkvBlock.__init__NFcCsz|jdkr||}|j||||d\}}||}|j|||d\}}||}||f}|rn||f7}n|d7}|S)Nr)rSrrr)rtrrrrr)rrrSroutput_attentionsrroutputsrrrrK}s   zRwkvBlock.forward)NFFrrrrrrns rc@s<eZdZdZeZdZdgZddgZdZ ddZ d d d Z d S)RwkvPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. rrrOrPTc st|trN|j}|jj}|jj|j||dd||}tjfddt D|j j |j j d}|ddddf}fddt D}tj||j j |j j d}tjddt D|jj |jj dd }th||j _t|jtd ||j_t|||j _t||d |j_t|d ||j_W5QRXnt|tr|j}|jj}|jjd||}tjfd dt D|j j |j j d}|ddddf}t&t|||j _t|||j_W5QRXdS) zInitialize the weights.rg?csg|] }|qSrrrirUrrrsz5RwkvPreTrainedModel._init_weights..r2r3Ncs,g|]$}dd|dddqS)rgffffff?g?r)rh)ru ratio_0_to_1rrrscSsg|]}|dddqS)rr rrrrrrsg?g333333?csg|] }|qSrrrrrrrs) isinstancerlrtrrnum_hidden_layersrUrur@Ztensorr\rxr2r3rOrPno_graddataZ ones_likemathlogpowryrzr)rmodulertrZratio_1_to_almost0Z time_weightZ decay_speedZzigzagr)rurUrr _init_weightssZ         z!RwkvPreTrainedModel._init_weightsFcCst|tr||_dSr)r RwkvModelgradient_checkpointing)rrrRrrr_set_gradient_checkpointings z/RwkvPreTrainedModel._set_gradient_checkpointingN)F) rXrYrZ__doc__r config_classZbase_model_prefixZ_no_split_modulesZ_keep_in_fp32_modulesZsupports_gradient_checkpointingrrrrrrrs9rc@sbeZdZUdZdZejed<dZe e ejed<dZ e e ejed<dZ e e ejed<dS) RwkvOutputa Class for the RWKV model outputs. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlast_hidden_staterS hidden_states attentions)rXrYrZrrr@ FloatTensor__annotations__rSrrrrrrrrrrs rc@steZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dZee ejed<dS)RwkvCausalLMOutputa| Base class for causal language model (or autoregressive) outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). state (list of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`): The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to avoid providing the old `input_ids`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. NlosslogitsrSrr)rXrYrZrrrr@rrrrSrrrrrrrrrs ra> This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`RwkvConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a Args: 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) attention_mask (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. This is currently not used by `RwkvModel`, but will be supported in the future. [What are attention masks?](../glossary#attention-mask) inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. state (tuple of five `torch.FloatTensor` of shape `(batch_size, hidden_size, num_hidden_layers)`, *optional*): If passed along, the model uses the previous state in all the blocks (which will give the output for the `input_ids` provided as if the model add `state_input_ids + input_ids` as context). use_cache (`bool`, *optional*): If set to `True`, the last state is returned and can be used to quickly generate the next logits. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. z^The bare RWKV Model transformer outputting raw hidden-states without any specific head on top.cseZdZfddZddZddZeeee e e dde e je e je e je ee je ee ee ee eeee fd d d Zd d ZddZZS)rcsdttjj|_tfddtj D|_ t j|_ d|_ d|_|dS)Ncsg|]}t|dqS)r)r)ridxrrrrr]sz&RwkvModel.__init__..F)rprqrZ Embedding vocab_sizerU embeddingsZ ModuleListr\rblocksrln_outlayers_are_rescaledr post_initrrrrrrrqYs  zRwkvModel.__init__cCs|jSrrrrrrget_input_embeddingsgszRwkvModel.get_input_embeddingscCs ||_dSrrrZnew_embeddingsrrrset_input_embeddingsjszRwkvModel.set_input_embeddings checkpoint output_typerN) input_idsattention_mask inputs_embedsrSrroutput_hidden_states return_dictreturnc sHdk r n|jj|dk r |n|jj}dk r4n|jsB|jjnd|dk rR|n|jj}|j|jkrn||dk rdk rtdn|dkrdkrtddkr| |r|dkr d|jj |jj ffddt dD}|dd 8<|jr"|jr"r"td d} r0d nd} |r>d nd} t|jD]\} } |jr|jrfd d }tjj|| | |\} }}n| | |d\} }}|jr|jjdkr| d|jjdkr| d} |r| | f} rL| |f} qL|| } |r| | f} |s8tdd| || | fDSt| || | dS)NFzDYou cannot specify both input_ids and inputs_embeds at the same timez5You have to specify either input_ids or inputs_embedsrcs0g|](}tj|dkrjntjjdqS)rr)r@rGr2rHr3r)rshaperrrs z%RwkvModel.forward..rgꌠ9Y>)FzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...rcsfdd}|S)Ncs|dS)N)rrr)inputs)rrrrrcustom_forwardszHRwkvModel.forward..create_custom_forward..custom_forwardr)rr)rr)rrcreate_custom_forwardsz0RwkvModel.forward..create_custom_forward)rSrrrr4css|]}|dk r|VqdSrr)rxrrrrgsz$RwkvModel.forward..)rrSrr)rrrrtrainingruse_return_dictr_rescale_layersr<rr;rUrr\rr'Z warning_once enumeraterr@utilsr rescale_everyrrr)rrrrrSrrrrrZall_self_attentionsZall_hidden_statesrblockrrr)rrrrrrKms           zRwkvModel.forwardc Csl|j|j krdS|jjdkr^t.t|jD]\}}|jr|jj j dt ||jj|j jj dt ||jjq6t|jj j dr|jj j jdt ||jj|j jj jdt ||jjq6t|jj j dr||jj |||j j|q6|jj j dt ||jj|j jj dt ||jjq6W5QRX|j |_dS)Nrr4SCB quant_state)rrrrrr@rrrrrVweightZmul_intrrRhasattrrdiv_ _bnb_4bit_dequantize_and_rescale)rblock_idrrrrrs"  ""$ ,zRwkvModel._rescale_layerscCsttstdddl}|j|jj|jj}|dt ||j j |j j |ddd|j}t|d|dS) z Perform the dequantization and rescaling of the weights of a given layer. After that operation the layer will be quantized again. z/Please install bitsandbytes to use this method.rNr4cpuF)Z requires_gradr)r ImportErrorZ bitsandbytesZ functionalZdequantize_4bitrrrrrrrrrZ Params4bitrMr3setattr)rZ target_layerrZbnbZdequant_weightsZ quant_weightrrrrsz*RwkvModel._bnb_4bit_dequantize_and_rescale)NNNNNNNN)rXrYrZrqrrrRWKV_INPUTS_DOCSTRINGr _CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOCrr@ LongTensorrrboolrrrKrrrrrrrrTs<   `rz The RWKV Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). cseZdZdgZfddZddZddZdd d Zee e e e e d deejeejeejeeejeejeeeeeeeeeee fd d dZZS)RwkvForCausalLMz head.weightcs8t|t||_tj|j|jdd|_| dS)NFro) rprqrrrr}rUrheadrrrrrrq s  zRwkvForCausalLM.__init__cCs|jSrrrrrrget_output_embeddingssz%RwkvForCausalLM.get_output_embeddingscCs ||_dSrrrrrrset_output_embeddingssz%RwkvForCausalLM.set_output_embeddingsNcKsL|dk r|dddfd}|dk r8|dkr8d|i}nd|i}||d<|S)NrnrrrSr7)rrrSrkwargsZ model_inputsrrrprepare_inputs_for_generations z-RwkvForCausalLM.prepare_inputs_for_generationr) rrrrSlabelsrrrrrc  Cs| dk r | n|jj} |j||||||| d} | d} || } d} |dk r|| j}| dddddf}|dddf}t}||d| d|d} | s| f| dd}| dk r| f|S|St | | | j | j | j dS)a 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)rrSrrrrr.rnr)rrrSrr)rrrrrrMr3rCr viewr;rrSrr)rrrrrSrrrrrZ rwkv_outputsrrrZ shift_logitsZ shift_labelsZloss_fctrVrrrrK's:   zRwkvForCausalLM.forward)NN) NNNNNNNNN)rXrYrZZ_tied_weights_keysrqrrrrrr rrrrr@rrrrrrrKrrrrrrs@    r)NF)NF)5rr dataclassesrpathlibrtypingrrrrr@Ztorch.utils.checkpointrZtorch.nnr Zmodeling_utilsr rr r rrrrrrZconfiguration_rwkvrZ get_loggerrXr'rrZ"RWKV_PRETRAINED_MODEL_ARCHIVE_LISTr%r,ZautogradFunctionr-rerkModulerlrrrrrZRWKV_START_DOCSTRINGrrrrrrrsf     (   j , F&"J +*