[TU]5 [US5(aZ[UR[ 5(a;URR SURR S55UlOSUlURUl URUl Xl [URUl URURU5uo0lURSUSS9 UR Ulg)N rope_scaling rope_typetypedefaultinv_freqF) persistent)super__init__hasattr isinstancer!dictgetr"max_position_embeddingsmax_seq_len_cachedoriginal_max_seq_lenrr rope_init_fnattention_scalingregister_bufferr%original_inv_freq)selfrdevicer% __class__s d/var/www/html/shao/venv/lib/python3.13/site-packages/transformers/models/cohere2/modeling_cohere2.pyr(Cohere2RotaryEmbedding.__init__+s  6> * *z&:M:Mt/T/T#0044[&BUBUBYBYZ`BabDN&DN"("@"@$*$B$B! /?+/+<+freqsembrMrNs r7forwardCohere2RotaryEmbedding.forward<sB!MM$4-8>>@GG HZHZ[\H]_acde ,QaZ 8 > > @'1!((--'E'E!((--[`J`ahhmmfk ^^ U C&,,.1F1L1L1NNYYZ[]^_E))%;C'')d444C'')d444C D vvAGGv$cff177f&;;; D Cs BF F)r1rr.r3r/r0r"N) __name__ __module__ __qualname____firstlineno__rr(rIno_gradrrV__static_attributes__ __classcell__r6s@r7rr*s6/}//" ]]_ < [TU]5 [R"[R "U55UlX lg)zcThe hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dimN)r'r(nn ParameterrIonesweightvariance_epsilon)r4 hidden_sizeepsbiasr6s r7r(Cohere2LayerNorm.__init__Ms- ll5::k#:; #r9cURnUR[R5nUR SSS9nX- R S5R SSS9nX- [R "X@R-5-nURR[R5U-nURU5$)Nr;T)keepdimr@) rDrOrIfloat32meanpowrsqrtrirh)r4 hidden_states input_dtyperqvariances r7rVCohere2LayerNorm.forwardSs#)) %((7 !!"d!3!(--a055b$5G&-XH]H]=]1^^  u}}5 E  ,,r9)rirh)Ngh㈵>FrYrZr[r\r(rVr^r_r`s@r7rbrbLs$ --r9rbrtn_repreturncURup#pEUS:XaU$USS2SS2SSS2SS24RX#XU5nURX#U-XE5$)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)rGrFreshape)rtrybatchnum_key_value_headsslenhead_dims r7 repeat_kvr]s_ 2?1D1D.E z!!Qa"23::5W\dlmM  e(CT TTr9modulequerykeyvalueattention_maskscalingdropoutkwargsc @[X R5n[X0R5n [R"XR SS55U-n Ub"USS2SS2SS2SUR S24n X-n [ RRU S[RS9RUR5n [ RRXURS9n [R"X5n U R SS5R5n X4$)Nr@rr;)rBrD)ptrainingr)rnum_key_value_groupsrImatmulrKrGre functionalsoftmaxrprOrDrr contiguous) rrrrrrrr key_states value_states attn_weights causal_mask attn_outputs r7eager_attention_forwardris3 ; ;C/0Gw;q>C/0G ::E: "GJJUJ$; ;;r9cP^\rSrSrSrSS\S\\4U4SjjjrSS\ RS\ \ R\ R4S\\ RS \\ S \\ RS \\S \ \ R\\ R\\ \ R44S jjrSrU=r$)Cohere2Attentionz=Multi-headed attention from 'Attention Is All You Need' paperr layer_idxc>[TU]5 XlX l[ USUR UR -5UlUR UR-Ul URS-Ul URUl SUl URUS:Xa UROSUl[R "UR UR UR-UR"S9Ul[R "UR URUR-UR"S9Ul[R "UR URUR-UR"S9Ul[R "UR UR-UR UR"S9Ulg)NrgTsliding_attentionrl)r'r(rrgetattrrjnum_attention_headsrr~rrattention_dropout is_causal layer_typessliding_windowreLinearattention_biasq_projk_projv_projo_projr4rrr6s r7r(Cohere2Attention.__init__sm  " F4F4F&JdJd4de $*$>$>&B\B\$\!}}d* !'!9!97=7I7I)7TXk7kf33quii    : :T]] JQWQfQf  ii    : :T]] JQWQfQf  ii    : :T]] JQWQfQf  ii  & & 68J8JQWQfQf  r9rtposition_embeddingsrpast_key_valuecache_positionrrzc dURSSn/UQSPURP7nURU5RU5R SS5n UR U5RU5R SS5n UR U5RU5R SS5n UupURb[XX5upUb$XUS.nURXURU5up[nURRS:wa[URRnU"UU U U U4UR(dSO UR UR"URS.UD6unnUR$"/UQSP76R'5nUR)U5nUU4$)Nr;rr@)rNrMreager)rrr)rGrrviewrKrrrrupdaterrr_attn_implementationrrrrr|rr)r4rtrrrrr input_shape hidden_shape query_statesrrrMrN cache_kwargsattention_interfacerrs r7rVCohere2Attention.forwards$))#2. 88b8$--8 {{=166|DNNqRST [[/44\BLLQPQR {{=166|DNNqRST &    *';LVY'_ $L  %#&nUL'5'<'*)!. *) ! *) !!1!12 *)-.*) u||Xell3XeELL>Q5RR S*)*)r9rc.^\rSrSrU4SjrSrSrU=r$) Cohere2MLPc>[TU]5 XlURUlURUl[ R "URURSS9Ul[ R "URURSS9Ul[ R "URURSS9Ul [URUl gNFr) r'r(rrjintermediate_sizerer gate_projup_proj down_projr hidden_actact_fnr4rr6s r7r(Cohere2MLP.__init__s  !--!'!9!94#3#3T5K5KRWXyy!1!143I3IPUV 4#9#94;K;KRWXV../ r9cURURURU55URU5-5nU$rX)rrrr)r4rPrs r7rVCohere2MLP.forwards6NN4;;t~~a/@#ADLLQRO#ST r9)rrrrrjrrrxr`s@r7rrs0r9rcL^\rSrSrS\S\4U4SjjrSS\RS\ \R\R4S\ \RS\ \ S \ \ S \ \RS \\S \ \R \ \ \R \R 444S jjrSrU=r$)Cohere2DecoderLayerirrc>[TU]5 URUl[XS9Ul[ U5Ul[URURS9Ul URUUl g)N)rrrjrk) r'r(rjr self_attnrmlprblayer_norm_epsinput_layernormrattention_typers r7r(Cohere2DecoderLayer.__init__sc !--)Mf%/V=O=OV\VkVkl$00;r9rtrrr use_cacherrrzc UnURU5nUR"SUUUUUUS.UD6upURU5n X-U -nU$)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_size, sequence_length)` if flash attention is used or `(batch_size, 1, query_sequence_length, key_sequence_length)` if default attention is used. past_key_value (`Tuple(torch.FloatTensor)`, *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 position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. )rtrrrrr)rrr) r4rtrrrrrrresidualhidden_states_attention_hidden_states_mlps r7rVCohere2DecoderLayer.forward sq:!,,]; %)^^& ' 3))) & & "!HH]3 :=NN r9)rrjrrr)NNFN)rYrZr[r\rrr(rIrrrrboolrrr FloatTensorrVr^r_r`s@r7rrs<}<<26*.$)59+||+#5<<#=>+!. + ! + D> +!!1!12+-.+ u  (51B1BEDUDU1U+V"WW X++r9rcR\rSrSr%\\S'SrSrS/rS/r Sr Sr Sr Sr Sr\\S.rSrg ) Cohere2PreTrainedModeli8rmodelTrpast_key_values)rt attentionsrN)rYrZr[r\r__annotations__base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_supports_flex_attn_can_compile_fullgraph_supports_attention_backendrr_can_record_outputsr^rr9r7rr8sQ &*#./#4"5N!"&,&r9rc^\rSrSrS\4U4Sjjr\\SS\\ RS\\ RS\\ RS\\ S\\ RS \\S \\ RS \\S \4S jj55rSrU=r$) Cohere2ModeliKrc >[TU]U5 URUlURUl[ R "URURUR5Ul[ R"[UR5Vs/sHn[X5PM sn5Ul [URURS9Ul[#US9UlSUlUR)5 gs snf)Nr)rF)r'r( pad_token_id padding_idx vocab_sizere Embeddingrj embed_tokens ModuleListrangenum_hidden_layersrlayersrbrnormr rotary_embgradient_checkpointing post_initrs r7r(Cohere2Model.__init__Ms  !.. ++LL):):F input_idsrrQr inputs_embedsrrrrzc USLUSL- (a [S5eUcURU5nU(aUcUR(d [5nUcDUbUR 5OSn [ R "XURS-URS9nUcURS5n[U=n [5(d*URUUUUUS.n [S 0U D6[S 0U D6S.n Un URX5n UR HnU"U 4U XR"UUUS.UD6n M UR%U 5n ['U US9$) Nz:You must specify exactly one of input_ids or inputs_embedsrr)r5)r input_embedsrrrrQ)full_attentionr)rrrrr)last_hidden_staterr) ValueErrorr rr get_seq_lengthrIarangerGr5rr*r+rr r rrrrr)r4rrrQrrrrrpast_seen_tokenscausal_mask_mapping mask_kwargsrtr decoder_layers r7rVCohere2Model.forward]sw  -t"; <YZ Z   --i8M 0*nO  !CRC^==?de "\\ ]5H5H5K"KTaThThN  )33A6L?-FF++ -"0"0#2 , K#5"C{"C%F%U%U#  & "oomJ![[M)$723O3OP.#- M) -0 &++  r9)r rrrrrr )NNNNNNN)rYrZr[r\rr(rrrrIrrrrrrrrrVr^r_r`s@r7rrKs} 151537+/59$(59< E,,-< !.< u//0 < "% <   1 12 < D>< !!1!12< +,<  !< < r9rc^\rSrSrS/rSS0rSS/S/40rU4SjrSrS r \ \ SS \ \ RS \ \ RS \ \ RS \ \\\\ R&4S\ \ R&S\ \ RS\ \S\ \S\ \S\ \ RS\\\ R4S\\S\4Sjj55rSrU=r$)Cohere2ForCausalLMizlm_head.weightlm_head colwise_reprtlogitsc(>[TU]U5 [U5UlURUl[ R "URURSS9UlURUl URUl UR5 gr) r'r(rrr rerrjr% logit_scaletie_word_embeddingsrrs r7r(Cohere2ForCausalLM.__init__sq  !&)  ++yy!3!3V5F5FUS !--#)#=#=  r9cXlgrXr)r4decoders r7 set_decoderCohere2ForCausalLM.set_decoders r9cUR$rXr-)r4s r7 get_decoderCohere2ForCausalLM.get_decoders zzr9rrrQrrlabelsroutput_attentionsoutput_hidden_statesrlogits_to_keeprrzc UbUOURRnU bU OURRn UR"SUUUUUUUU U S. U D6n U Rn[ U [ 5(a [U *S5OU nURUSS2USS245nUUR-nSnUb)UR"SUX`RRS.U D6n[UUU RU RU RS9$)a 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, Cohere2ForCausalLM >> model = Cohere2ForCausalLM.from_pretrained("Cohere2ForAI/c4ai-command-r-v01") >> tokenizer = AutoTokenizer.from_pretrained("Cohere2ForAI/c4ai-command-r-v01") >> 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) rrrQrrrr5r6r)r'r4r )lossr'rrtrr)rr5r6rrr*rslicer%r) loss_functionr rrrtr)r4rrrQrrr4rr5r6rr7routputsrt slice_indicesr'r9s r7rVCohere2ForCausalLM.forwards(N2C1N-TXT_T_TqTq$8$D $++JjJj  ,0:: , )%+'/!5) ,  ,  11 8B>SV8W8W~ot4]k mA}a,?@A$***  %%pVF{{OeOepiopD%#33!//))   r9)r%r)rr*r ) NNNNNNNNNNr)rYrZr[r\_tied_weights_keys_tp_plan_pp_planr(r/r2rrrrIrrrrlistrrrrrrrVr^r_r`s@r7r$r$s*+=)H_-z:;H 151537KO59-1$(,0/35934H E,,-H !.H u//0 H "%tE4E4E/F(F"GH H   1 12 H ))*H D>H $D>H 'tnH !!1!12H c5<</0H +,H  H H r9r$)r$rr)r)Nr)9typingrrrrItorch.nnre activationsr cache_utilsrr generationr masking_utilsr r modeling_flash_attention_utilsr modeling_layersrmodeling_outputsrrmodeling_rope_utilsrrmodeling_utilsrrprocessing_utilsrutilsrrr utils.genericrconfiguration_cohere2rModulerrbrrrrErrrrrrrrr$__all__rr9r7rTs,-, !.)RB9OKF&II/0