gdZddlZddlmZddlmZddlmZmZm Z m Z ddl Z ddl Z ddl m Z ddlmZdd lmZdd lmZdd lmZmZmZmZmZmZmZmZd d lmZeje Z!dZ"dZ#da$dZ%Gdde j&j'Z(d*dZ)d*dZ*Gdde j+Z,Gdde j+Z-Gdde j+Z.GddeZ/eGddeZ0eGd d!eZ1d"Z2d#Z3ed$e2Gd%d&e/Z4ed'e2Gd(d)e/eZ5dS)+zPyTorch RWKV model.N) dataclass)Path)ListOptionalTupleUnion)nn)CrossEntropyLoss)GenerationMixin)PreTrainedModel) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardis_bitsandbytes_availableis_ninja_availableis_torch_cuda_availablelogging) RwkvConfigzRWKV/rwkv-4-169m-pilercddlm}ttjjjdz dz fddD}t t j|krdSt d|dd d d d d dd|g}|d||tj tj k|a|t _dS)Nr)loadkernelsrwkvcg|]}|z Sr).0f kernel_folders b/var/www/html/ai-engine/env/lib/python3.11/site-packages/transformers/models/rwkv/modeling_rwkv.py z(load_wkv_cuda_kernel..:sfffq*fff)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=wkv_)namesourcesverboseextra_cuda_cflags) torch.utils.cpp_extensionrr__file__resolveparentrwkv_cuda_kernelmax_seq_lengthloggerinfor get_verbosityDEBUG)context_length load_kernelcuda_kernel_filesflagsr s @r!load_wkv_cuda_kernelr84s ======NN**,,3:AIMPVVMffff4efff#(8(G>(Y(Y KKV^VVVWWW  &"."" E#{ $N $ $!&((GM9  '5###r#c>eZdZeddZeddZdS)RwkvLinearAttentionNFc B|\}}} |tjkr td|dtjd|| zt | dzdkr'td|d| dt | dd |j|_|jjd ks0|jjd ks |jjd ks|jjd krtd tj |  }|jtj kr<| }| }| }| }| }| }tj|tj } |s||Ktj|| d tj|jtj}|dddddfxxdzcc<n2tjd|Dd }|jtjkr tj} n tj} | ||||| |n<|jtjkr tjn tj} | ||||| |||||| |!dtj|d dD}| |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 r$cudazUCalling the CUDA kernel for wkv attention requires all tensors to be on CUDA devices. memory_formatr )dtypedevicer?籡*Gc8g|]}|dSrB unsqueezerss r!r"z/RwkvLinearAttention.forward..s""A"A"Aa1;;q>>"A"A"Ar#)dimc8g|]}|dSrE)squeezerHs r!r"z/RwkvLinearAttention.forward..s"HHHaQYYq\\HHHr#)sizer.r/ ValueErrorminr@ input_dtyperAtypetorchexpfloat contiguousfloat16 empty_likecontiguous_formatzerosfloat32catbfloat16forward_with_state_bf16forward_with_state forward_bf16forwardsave_for_backwardchunkto) ctx time_decay time_firstkeyvaluestate return_state batch_sizeseq_len hidden_sizeoutput forward_funcs r!r`zRwkvLinearAttention.forwardUsG+.88::( G[ %4 4 4FwFF#2FFF   #c+r&:&: :a ? ?7j77[77";33777  )   "f , , %//z&((| F**tuu ui 0 0 2 2 = = ? ?@@@ 9 % %#))++J))++CKKMME**,, nn  ""!#U5LMMM  E5,} -:"'"9 aaaAg$& "A"A5"A"A"AqIIITTVVyEN**/G /B LZeVU K K K K)rP saved_tensorsrRrWrXr\rZrVrTr. backward_bf16backwardrUrc)rdg_outputg_staterPrerfrgrhrn g_time_decay g_time_firstg_keyg_value backward_funcs r!rszRwkvLinearAttention.backwardseo 585F2 JUF' 1$/5>$A$A%..u}   ' %BYZZZ  E4KLLL"58OPPP %- ' '~~''H:E:W:W(66]m]v          ! !      OOK ( ( OOK ( ( HH[ ! ! JJ{ # #     r#NFN)__name__ __module__ __qualname__ staticmethodr`rsrr#r!r:r:TsS<1<1<1\<1|% % % \% % % r#r:Fc|\}}}tj|}|tj|dddftj} tj|dddftj} tj|dddftjdz } n|\} } } tj| }t |D] } |dd| f} |dd| f}tj| | |z}tj| |z }tj| |z|z }|| z||zz}|| z|z}||z |j |dd| f<tj| |z| }tj| |z|z }tj| |z }|| z||zz} || z|z} |}  |s|| | | g}||fS)Nr)r@rC) rMrR zeros_likerZrSrangerTmaximumrcr@)rerfrgrhrirj_ seq_lengthrn num_state den_state max_state current_index current_key current_valuemax_for_outpute1e2 numerator denominator max_for_states r!rwkv_linear_attention_cpurs$xxzzAz1  c " "F }$SAYemDDD $SAYemDDD $SAYemDDDtK */' 9i )J'''Jz**"" !!!]*+1133 aaa./ y+ 2JKK Yy>1 2 2 Y{Z/.@ A ANR-%77 9nr) $- $;#?#? #M#Mqqq-  i*& > Y{]2 3 3NR-%77 NR' ! 2u(Iy1 5=r#ctd||||fD}|ddk}t|s|rt||||||St||||||S)Nc36K|]}|jjdkVdS)r=N)rArQ)rts r! z(rwkv_linear_attention..s+XXa!(-6)XXXXXXr#rrirj)anyrMr.rr:apply)rerfrgrhrirjno_cuda one_tokens r!rwkv_linear_attentionrsXXJ CQV3WXXXXXG q I7i(ZeSXgstttt"((ZeUT`aaar#c0eZdZdfd ZddZd dZxZS) RwkvSelfAttentionrcdt||_tduotj|jk}t rPtrB|s@ t|jn*#t$rt dYnwxYw||_ |j }|j|jn|}||_tjt#j||_tjt#j||_tjt#jdd||_tjt#jdd||_tjt#jdd||_tjd|_tj||d|_tj||d|_tj||d|_tj||d|_dS)Nz9Could not load the custom CUDA kernel for RWKV attention.rrrrFbias)super__init__configr.r/r4rrr8 Exceptionr0r1layer_idrmattention_hidden_sizer ParameterrRemptyrerf time_mix_keytime_mix_valuetime_mix_receptance ZeroPad2d time_shiftLinearrgrh receptancern)selfrr kernel_loadedrmr __class__s r!rzRwkvSelfAttention.__init__s  (4q9I9X\b\q9q    Y$;$=$= Ym Y Y$V%:;;;; Y Y Y WXXXXX Y  ( ,2,H,TF ( (Ze &;",u{3H'I'IJJ,u{3H'I'IJJLQ;)G)GHH l5;q![+I+IJJ#%< Aq+0N0N#O#O ,}559[*?eLLLY{,ANNN )K1FUSSSi 5{OOO s&A;;$B"!B"Ncb|ddkr||ddddd|jf}n8||}|!|ddddd|jf|dddf<||jz|d|jz zz}||jz|d|jz zz}||jz|d|jz zz}||}||}tj | |}|!|dddf|ddddd|jf<||||fSNrrr) rMrrrrrrgrhrRsigmoidr)rhiddenrishiftedrgrhrs r!extract_key_valuez#RwkvSelfAttention.extract_key_valuesW ;;q>>Q  5#4Ahqqq!!!T]23GGoof--G %aAAAt})< =1 t((7a$:K6K+LL,,w!d>Q:Q/RRd66AH`D`9aa hhsmm 5!!]4??:#>#>??  ,2111b5ME!HQQQ4=( )3u,,r#Fc||\}}}}|#tfd|ddDnd}tjj||||\}}|W|d|dddddjf<|d|dddddjf<|d|dddddjf<||z|fS) Nric3@K|]}|ddddjfVdSr|r)rrIrs r!rz,RwkvSelfAttention.forward..)s9FFqAaaaDM12FFFFFFr#rBrrrr )rtuplerrerfrrn) rrri use_cacherrgrh layer_staters ` r!r`zRwkvSelfAttention.forward's(,(>(>vU(>(S(S% CJOJ[eFFFFE!""IFFFFFFae 1 O O  "    k  ",7NE!HQQQ4=( ),7NE!HQQQ4=( ),7NE!HQQQ4=( ){{:,--u44r#rr|r{)r}r~rrrr` __classcell__rs@r!rrskPPPPPP<----&55555555r#rc(eZdZdfd ZddZxZS)RwkvFeedForwardrc0t||_||_|j}|j|jn d|jz}t jd|_t j tj dd||_ t j tj dd||_ t j||d|_t j||d|_t j||d|_dS)NrrrFr)rrrrrmintermediate_sizer rrrrRrrrrrgrrh)rrrrmrrs r!rzRwkvFeedForward.__init__<s    ( (.(@(LF $ $RSV\VhRh ,}55LQ;)G)GHH#%< Aq+0N0N#O#O 9[*;%HHH)K5IIIY0+EJJJ r#Nc||ddkr||ddddd|jf}n8||}|!|ddddd|jf|dddf<||jz|d|jz zz}||jz|d|jz zz}t jt j||}| |}t j | |}|!|dddf|ddddd|jf<||z|fSr) rMrrrrrRsquarerelurgrhrr)rrrirrgrrhs r!r`zRwkvFeedForward.forwardMsK ;;q>>Q  5#4Ahqqq!!!T]23GGoof--G %aAAAt})< =1 t((7a$:K6K+LLd66AH`D`9aa l5:dhhsmm4455 3]4??:#>#>??  ,2111b5ME!HQQQ4=( )E!5((r#rr|r}r~rrr`rrs@r!rr;sWKKKKKK"))))))))r#rc&eZdZfdZddZxZS) RwkvBlockct||_||_|dkr%t j|j|j|_t j|j|j|_ t j|j|j|_ t|||_ t|||_dS)Nr)eps)rrrrr LayerNormrmlayer_norm_epsilonpre_lnln1ln2r attentionr feed_forward)rrrrs r!rzRwkvBlock.__init__bs    q==,v'9v?XYYYDK< 28QRRR< 28QRRR*68<<+FH==r#NFc4|jdkr||}|||||\}}||z}||||\}}||z}||f}|r||fz }n|dz }|S)Nr)rirrr|)rrrrrr)rrriroutput_attentionsrroutputss r!r`zRwkvBlock.forwardps =A  [[((F>>$((6*:*:%S\>]] 5)#"//0@0@/NN e,&5/    | #GG w Gr#)NFFrrs@r!rrasL > > > > >r#rc6eZdZdZeZdZdgZddgZdZ dZ dZ dS) RwkvPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. rrrerfTc t|tr|j}|jj}|jj |j||dz z d||z z }tj fdt D|j j |j j }|ddddf} fdtD}tj||j j |j j }tjdtD|jj |jj dz}tj5||j _tj|jt%jd z|z|j_tj|||j _tj||d zz|j_tj|d|z|j_ddddS#1swxYwYdSt|t.r|j}|jj}|jj d||z z }tj fd t D|j j |j j }|ddddf}tj5tj|||j _tj|||j_ddddS#1swxYwYdSdS) zInitialize the weights.rg?cg|]}|z Srrrirms r!r"z5RwkvPreTrainedModel._init_weights..===Q[===r#r@rANc>g|]}dd|dz z ddzzzzzS)rgffffff?g?r)rhr ratio_0_to_1s r!r"z5RwkvPreTrainedModel._init_weights..sMQ!4q89sS.s$KKKa!eq[1_KKKr#g?g333333?cg|]}|z Srrrs r!r"z5RwkvPreTrainedModel._init_weights..rr#) isinstancerrrnum_hidden_layersrmrrRtensorrrr@rArerfno_graddata ones_likemathlogpowrrr) rmodulerrratio_1_to_almost0 time_weight decay_speedzigzagrrmrs @@@r! _init_weightsz!RwkvPreTrainedModel._init_weightss f/ 0 05 ]H & ?  -3K$*$@ !#'81'<=L!$3D(D!E ,====% *<*<===)/*1K &dD!!!m4K455K ,{&:K:QZ`ZkZrsssK KKe4I.J.JKKK +1!,3    c c)4!&).9JTXVY]]9Z]c9c)d)d!&+09[BT+U+U#(-2Y{DV-W-WZ]`lZl-l%*27)KOaIa2b2b*/  c c c c c c c c c c c c c c c c c c 0 0 ]H & ?  -3K!$3D(D!E ,====% *<*<===)/*1K &dD!!!m4K ] ]+09[BT+U+U#(27)KI[2\2\*/ ] ] ] ] ] ] ] ] ] ] ] ] ] ] ] ] ] ] ] ]s%B.G;;G?G?(?K44K8;K8N) r}r~r__doc__r config_classbase_model_prefix_no_split_modules_keep_in_fp32_modulessupports_gradient_checkpointing _is_statefulrrr#r!rrs] L$ )<8&*#L7]7]7]7]7]r#rceZdZUdZdZejed<dZe e ejed<dZ e e ejdfed<dZ e e ejdfed<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_stateri. hidden_states attentions)r}r~rrrrR FloatTensor__annotations__rirrrrrrr#r!rrs,,0u(////3E8D*+ ,333=AM8E%"3S"89:AAA:>Ju0#567>>>>>r#rceZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejdfed<dZee ejdfed<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. Nlosslogitsri.rr)r}r~rrr rrRrr r rirrrrrr#r!r r s0)-D(5$ %,,, $FE $$$/3E8D*+ ,333=AM8E%"3S"89:AAA:>Ju0#567>>>>>r#r a> 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.cbeZdZfdZdZdZeeee e e  dde e jde e jde e jd e ee jd e ed e ed e ed e edeee ffdZdZdZxZS) RwkvModelcttjjj|_tjfdtj D|_ tj j|_ d|_ d|_|dS)Nc2g|]}t|S)r)r)ridxrs r!r"z&RwkvModel.__init__..Ms&$p$p$pYv%D%D%D$p$p$pr#F)rrr Embedding vocab_sizerm embeddings ModuleListrrblocksrln_outlayers_are_rescaledgradient_checkpointing post_initrrrs `r!rzRwkvModel.__init__Is    ,v'8&:LMMm$p$p$p$pPUV\VnPoPo$p$p$pqq l6#566 #( &+# r#c|jSr|rrs r!get_input_embeddingszRwkvModel.get_input_embeddingsWs r#c||_dSr|rrnew_embeddingss r!set_input_embeddingszRwkvModel.set_input_embeddingsZs (r# checkpoint output_typerN input_idsattention_mask inputs_embedsrirroutput_hidden_states return_dictreturnc ||n |jj}||n |jj}||n|js |jjnd}||n |jj}|t d|j|jkr| |td|td| ||rZ|X d|jj |jjffdtdD}|dxxd zcc<|jr%|jr|rt d d}} |rd nd} |rd nd} t#|jD]\} } |jr*|jr#|| j| |||\} }}n| | ||| \} }}|jr+|jjdkr| d z|jjzdkr| dz } |r| | fz} |r| |fz} || } |r| | fz} |st/d| || | fDSt1| || | S)NFz<`attention_mask` was passed, but it is unused in this model.zDYou cannot specify both input_ids and inputs_embeds at the same timez5You have to specify either input_ids or inputs_embedsrclg|]0}tj|dkrjn tjjd1S)rr)rRrYr@rZrA)rrr*shapes r!r"z%RwkvModel.forward..sY a-"5"5U][h[or#rgꌠ9Y>)FzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...r)rirrrrBc3K|]}||V dSr|r)rxs r!rz$RwkvModel.forward..s(ttqfgfsfsfsfsfsttr#)rrirr)rrr+trainingruse_return_dictr0 warning_oncer_rescale_layersrNrrMrmrrr enumerater_gradient_checkpointing_func__call__ rescale_everyrrr)rr(r)r*rirrr+r,rall_self_attentionsall_hidden_statesrblockrr0s ` @r!r`zRwkvModel.forward]sF"2C1N--TXT_Tq$8$D $+Jj "+!6IIZ^Zg=rT[=R=Rmr %0%cdd d  =#8TUU U   OOI66M  "''**DK,CT[EbcEq E !HHH HHH  & "4= " "##p" % $5?bb4"6@BBD#DK00 J JJC* t} 373T3TNM5)EV440 ujj495!)Wh4440 uj ( 2K-111W 99Q>> - 1 # I$58H$H!  J&9ZM&I# M22  E 1]4D D  utt]E;LNa$bttttt t++*     r#c ,|j|j krdS|jjdkrbt j5t |jD] \}}|jr|jj j dt||jjzz|j jj dt||jjzzt|jj j dr|jj j jdt||jjzz|j jj jdt||jjzz=t|jj j drB||jj |||j j||jj j dt||jjzz|j jj dt||jjzz" dddn #1swxYwY|j |_dS)NrrBSCB quant_state)rr4rr;rRrr8rrrnweightmul_intrrhhasattrr@div_ _bnb_4bit_dequantize_and_rescale)rblock_idr>s r!r7zRwkvModel._rescale_layerss  #DM(9 : : F ; $q ( ( r r'0'='=rrOHe} r.5::1HPTP[PiDi@j@j;jkkk*07<>qCTXT_TmHmDnDn?nooo!.4;@@c(VZVaVoJoFpFpApqqqqr r r r r r r r r r r r r r r r"(,}#4   sH7I<<JJctstdddl}|j|jj|jj}|dt||j j zz|j |dd|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.rNrBcpuF) requires_gradrB)r ImportError bitsandbytes functionaldequantize_4bitrBrrArFrDrr;r Params4bitrcrAsetattr)r target_layerrHbnbdequant_weights quant_weights r!rGz*RwkvModel._bnb_4bit_dequantize_and_rescales )** QOPP P"""".889L9QS_SfSrssQ#h$+2K&K"L"LLMMMv((););E)B)BRW(XX[[\k\rss  h 55555r#)NNNNNNNN)r}r~rrr r$rRWKV_INPUTS_DOCSTRINGr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOCrrR LongTensorrrboolrrr`r7rGrrs@r!rrDs     )))+*+@AA&$15595937$(,0/3&*Y Y E,-Y !!12Y   12 Y U./0 Y D> Y $D>Y 'tnY d^Y  uj !Y Y Y BA Y v55506666666r#rz The RWKV Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). ceZdZdgZfdZdZdZddZee e e e e  ddeejd eejd eejd eeejd eejd eedeedeedeedeee ffdZxZS)RwkvForCausalLMz head.weightct|t||_t j|j|jd|_| dS)NFr) rrrrr rrmrheadrrs r!rzRwkvForCausalLM.__init__s`    f%% If0&2C%PPP  r#c|jSr|r^rs r!get_output_embeddingsz%RwkvForCausalLM.get_output_embeddingss yr#c||_dSr|r`r"s r!set_output_embeddingsz%RwkvForCausalLM.set_output_embeddingss " r#Nc v||dddfd}||d|i}nd|i}||d<||d<|S)Nrr*r(rirrF)rr(rir*rkwargs model_inputss r!prepare_inputs_for_generationz-RwkvForCausalLM.prepare_inputs_for_generationsf  !!!!R%(22266I  $+];LL'3L % W$- [!r#r%r(r)r*rilabelsrrr+r,r-c p| | n |jj} |||||||| } | d} || } d} ||| j}| dddddf}|dddf}t}||d| d|d} | s| f| ddz}| | f|zn|St| | | j | j | j S)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)r*rirrr+r,r.rr)r r rirr)rr5rr^rcrArUr viewrMr rirr)rr(r)r*rirhrrr+r, rwkv_outputsrr r  shift_logits shift_labelsloss_fctrns r!r`zRwkvForCausalLM.forwardsl0&1%6 $D>6 'tn6 d^6  u(( )6 6 6 BA 6 6 6 6 6 r#r\r{)6rr dataclassesrpathlibrtypingrrrrrRtorch.utils.checkpointr torch.nnr generationr modeling_utilsr utilsrrrrrrrrconfiguration_rwkvr get_loggerr}r0rWrXr.r8autogradFunctionr:rrModulerrrrrr RWKV_START_DOCSTRINGrVrr\rr#r!r~s  !!!!!!//////////// %%%%%%))))))------                    +*****  H % %-555@g g g g g %.1g g g T))))XbbbbC5C5C5C5C5 C5C5C5L#)#)#)#)#)bi#)#)#)L DD]D]D]D]D]/D]D]D]N ??????? ?: ??????? ?@ (Vda6a6a6a6a6#a6a6 a6H ^ ^ ^ ^ ^ )?^ ^ ^ ^ ^ r#