+
    œ°³j�Æ  ã                   óþ  € ^ RI Ht ^ RIHt ^ RIt^ RIHt ^RIHt	 ^RI
Ht ^RIHtHt ^RIHt ^RIHtHt ^R	IHt ^R
IHt ^RIHtHt ^RIHtHt ^RIHtH t  ^RI!H"t" ^RI#H$t$H%t%H&t&H't'H(t( ^RI)H*t* ^RI+H,t,H-t- ^RI.H/t/ ^RI0H1t1 ^RI2H3t3 ^RI4H5t5H6t6 ^RI7H8t8 ](Pr                  ! ]:4      t;]&! RR7      ] ! R R]$4      4       4       t<]! R4       ! R R]Pz                  4      4       t> ! R R]Pz                  4      t? ! R  R!]Pz                  4      t@R" tA]! R#4      R>R$ l4       tBR% R& ltCR?R' R( lltD]! ]B4       ! R) R*]Pz                  4      4       tE ! R+ R,]4      tF]&! R-R7      ]& ! R. R/] 4      4       4       tG]& ! R0 R1]G4      4       tH ! R2 R3]Pz                  4      tI]&! R4R7       ! R5 R6]G]4      4       tJ ! R7 R8]Pz                  4      tK]& ! R9 R:]G4      4       tL]&! R;R7       ! R< R=]G]84      4       tM. R@OtNR# )Aé    )ÚCallable)Ú	dataclassN)Úinitialization)ÚACT2FN)ÚCacheÚDynamicCache)ÚGenerationMixin)Úuse_kernel_forward_from_hubÚuse_kernelized_func)Úcreate_causal_mask)ÚGradientCheckpointingLayer)ÚBaseModelOutputWithPastÚCausalLMOutputWithPast)ÚROPE_INIT_FUNCTIONSÚdynamic_rope_update)ÚALL_ATTENTION_FUNCTIONSÚPreTrainedModel)ÚUnpack)ÚModelOutputÚTransformersKwargsÚauto_docstringÚcan_return_tupleÚlogging)Údeprecate_kwarg)Úmaybe_autocastÚmerge_with_config_defaults)Úis_torchdynamo_compiling)Úcapture_outputs)Ú	AutoModel)Ú	CsmConfigÚCsmDepthDecoderConfig)ÚCsmGenerationMixinz:
    Base class for the model autoregressive outputs.
    )Úcustom_introc                   ó\   a € ] tR t^3t o RtRtRtRtRtRt	Rt
RtRtRtRtRtV 3R ltRtV tR# )ÚCsmOutputWithPasta=	  
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).
past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
    It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

    Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
    `past_key_values` input) to speed up sequential decoding.
depth_decoder_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
    Language modeling loss (for next-token prediction) of the depth decoder model.
depth_decoder_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
    Prediction scores of the depth decoder (scores for each vocabulary token before SoftMax).
depth_decoder_past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
    It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
depth_decoder_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.
depth_decoder_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)`.
backbone_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
    Language modeling loss (for next-token prediction) of the backbone model.
Nc                ó<  <€ V ^8„  d   Qh/ S[ P                  R,          ;R&   S[ P                  R,          ;R&   S[R,          ;R&   S[S[ P                  R3,          R,          ;R&   S[S[ P                  R3,          R,          ;R&   S[ P                  R,          ;R&   S[ P                  R,          ;R	&   S[R,          ;R
&   S[S[ P                  R3,          R,          ;R&   S[S[ P                  R3,          R,          ;R&   S[ P                  R,          ;R&   # )é   NÚlossÚlogitsÚpast_key_values.Úhidden_statesÚ
attentionsÚdepth_decoder_lossÚdepth_decoder_logitsÚdepth_decoder_past_key_valuesÚdepth_decoder_hidden_statesÚdepth_decoder_attentionsÚbackbone_loss)ÚtorchÚFloatTensorr   Útuple)ÚformatÚ__classdict__s   "€Úe/home/luca/screen-checking/.venv/lib/python3.14/site-packages/transformers/models/csm/modeling_csm.pyÚ__annotate__ÚCsmOutputWithPast.__annotate__3   s=  ø‡ ‚ ñF ×
Ñ
˜dÕ
"Ñ)ñG ñH ×Ñ Õ$Ñ+ñI ñJ ˜T•\Ñ(ñK ñL ™×*Ñ*¨CÐ/Õ0°4Õ7Ñ>ñM ñN ‘e×'Ñ'¨Ð,Õ-°Õ4Ñ;ñO ñP ×)Ñ)¨DÕ0Ñ7ñQ ñR  ×+Ñ+¨dÕ2Ñ9ñS ñT $)¨4¥<Ñ6ñU ñV "'¡u×'8Ñ'8¸#Ð'=Õ!>ÀÕ!EÑLñW ñX $¡E×$5Ñ$5°sÐ$:Õ;¸dÕBÑIñY ñZ ×$Ñ$ tÕ+Ñ2ò[ ó    © )Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__Ú__doc__r(   r)   r*   r+   r,   r-   r.   r/   r0   r1   r2   Ú__annotate_func__Ú__static_attributes__Ú__classdictcell__)r7   s   @r8   r%   r%   3   sP   ø‡ € ñð8 &*€DØ'+€FØ$(€OØ:>€MØ7;€JØ37ÐØ59ÐØ26Ð!ØHLÐØEIÐØ.2€M÷[ ƒ r;   r%   ÚRMSNormc                   óZ   a a€ ] tR t^ct oRV3R lV 3R llltV3R lR ltR tRtVtV ;t	# )Ú
CsmRMSNormc                ó$   <€ V ^8„  d   QhRS[ RR/# )r'   ÚepsÚreturnN)Úfloat)r6   r7   s   "€r8   r9   ÚCsmRMSNorm.__annotate__e   s   ø€ ÷ $ñ $©ð $¸$ñ $r;   c                óŽ   <€ \         SV `  4        \        P                  ! \        P
                  ! V4      4      V n        W n        R# )z)
CsmRMSNorm is equivalent to T5LayerNorm
N)ÚsuperÚ__init__ÚnnÚ	Parameterr3   ÚonesÚweightÚvariance_epsilon)ÚselfÚhidden_sizerI   Ú	__class__s   &&&€r8   rO   ÚCsmRMSNorm.__init__e   s/   ø€ ô 	‰ÑÔÜ—l’l¤5§:¢:¨kÓ#:Ó;ˆŒØ #Ör;   c                óN   <€ V ^8„  d   QhRS[ P                  RS[ P                  /# )r'   r+   rJ   ©r3   ÚTensor)r6   r7   s   "€r8   r9   rL   m   s#   ø€ ÷ ;ñ ;¡U§\¡\ð ;±e·l±lñ ;r;   c                ó:  € VP                   pVP                  \        P                  4      pVP	                  ^4      P                  RRR7      pV\        P                  ! W0P                  ,           4      ,          pV P                  VP                  V4      ,          # )r'   T)Úkeepdiméÿÿÿÿ)	ÚdtypeÚtor3   Úfloat32ÚpowÚmeanÚrsqrtrT   rS   )rU   r+   Úinput_dtypeÚvariances   &&  r8   ÚforwardÚCsmRMSNorm.forwardm   sw   € Ø#×)Ñ)ˆØ%×(Ñ(¬¯©Ó7ˆØ ×$Ñ$ QÓ'×,Ñ,¨R¸Ð,Ó>ˆØ%¬¯ª°H×?TÑ?TÕ4TÓ(UÕUˆØ�{‰{˜]×-Ñ-¨kÓ:Õ:Ð:r;   c                ó^   € \        V P                  P                  4       R V P                   2# )z, eps=)r5   rS   ÚshaperT   ©rU   s   &r8   Ú
extra_reprÚCsmRMSNorm.extra_reprt   s*   € Ü˜Ÿ™×)Ñ)Ó*Ð+¨6°$×2GÑ2GÐ1HÐIÐIr;   )rT   rS   )g�íµ ÷Æ°>)
r=   r>   r?   r@   rO   rg   rl   rC   rD   Ú__classcell__©rW   r7   s   @@r8   rG   rG   c   s%   ù‡ € ÷$õ $÷;ð ;÷Jò Jr;   rG   c                   óÎ   a a€ ] tR t^xt o]! RRR7      R
V3R lV 3R lll4       t]]! RRR7      R
V3R lR ll4       4       t]P                  ! 4       ]
R 4       4       tR	tVtV ;t# )ÚCsmRotaryEmbeddingÚdevicez5.18)Úversionc                ó    <€ V ^8„  d   QhRS[ /# ©r'   Úconfig©r    )r6   r7   s   "€r8   r9   ÚCsmRotaryEmbedding.__annotate__z   s   ø€ ÷ Oñ O™yñ Or;   c                óä  <€ \         SV `  4        VP                  V n        VP                  V n        Wn        V P
                  P                  R ,          V n        V P                  pV P                  R8w  d   \        V P                  ,          pV! V P
                  V4      w  q@n
        \        P                  ! VRR7      V n        \        P                  ! VP                  4       RR7      V n        R# )Ú	rope_typeÚdefaultF©Ú
persistentN)rN   rO   Úmax_position_embeddingsÚmax_seq_len_cachedÚoriginal_max_seq_lenrv   Úrope_parametersrz   Úcompute_default_rope_parametersr   Úattention_scalingrP   ÚBufferÚinv_freqÚcloneÚoriginal_inv_freq)rU   rv   rr   Úrope_init_fnr…   rW   s   &&&  €r8   rO   ÚCsmRotaryEmbedding.__init__y   s¬   ø€ ä‰ÑÔØ"(×"@Ñ"@ˆÔØ$*×$BÑ$BˆÔ!àŒàŸ™×4Ñ4°[ÕAˆŒØ!%×!EÑ!EˆØ�>‰>˜YÔ&Ü.¨t¯~©~Õ>ˆLÙ+7¸¿¹ÀVÓ+LÑ(ˆÔ(äŸ	š	 (°uÔ=ˆŒÜ!#§¢¨8¯>©>Ó+;ÈÔ!NˆÖr;   c                óP   <€ V ^8„  d   QhRS[ RS[S[P                  S[3,          /# )r'   rv   rJ   )r    r5   r3   r[   rK   )r6   r7   s   "€r8   r9   rx   Œ   s.   ø€ ÷ 5ñ 5±	ð 5ÑUZÑ[`×[gÑ[gÑinÐ[nÕUoñ 5r;   c           	     ó0  € V P                   R,          p\        V RR4      ;'       g    V P                  V P                  ,          pRpRV\        P
                  ! ^ V^\        P                  R7      V,          ,          ,          pVP                  V4      V3# )a~  
Computes the inverse frequencies according to the original RoPE implementation
Args:
    config ([`~transformers.PreTrainedConfig`]):
        The model configuration.
Returns:
    Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
    post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
Ú
rope_thetaÚhead_dimNg      ð?©r_   )r�   ÚgetattrrV   Únum_attention_headsr3   ÚarangerK   r`   )rv   rr   ÚkwargsÚbaseÚdimÚattention_factorr…   s   &&,    r8   r‚   Ú2CsmRotaryEmbedding.compute_default_rope_parametersŠ   s~   € ð ×%Ñ% lÕ3ˆÜ�f˜j¨$Ó/×cÐc°6×3EÑ3EÈ×IcÑIcÕ3cˆàÐà˜$¤5§<¢<°°3¸ÄÇÁÔ#MÐPSÕ#SÕTÕUˆØ�{‰{˜6Ó"Ð$4Ð4Ð4r;   c                ó^  € V P                   R,          P                  VP                  ^ ,          R	^4      P                  \        P
                  VP                  R7      pVR
,          P                  4       p\        VP                  P                  \        4      '       d2   VP                  P                  R8w  d   VP                  P                  MRp\        VRR7      ;_uu_ 4        W4,          P                  ^^4      p\        P                  ! Wf3R	R7      pVP                  4       V P                  ,          pVP                  4       V P                  ,          p	R R R 4       XP                  VP                   R7      X	P                  VP                   R7      3#   + '       g   i     LH; i)N)r_   rr   ÚmpsÚcpuF)Údevice_typeÚenabled©r”   rŽ   )NºNNNNr^   )r�   Nr�   )r…   Úexpandrj   r`   r3   rK   rr   Ú
isinstanceÚtypeÚstrr   Ú	transposeÚcatÚcosrƒ   Úsinr_   )
rU   ÚxÚposition_idsÚinv_freq_expandedÚposition_ids_expandedrš   ÚfreqsÚembr¤   r¥   s
   &&&       r8   rg   ÚCsmRotaryEmbedding.forwardž   s2  € ð �M‰M˜-Õ(×/Ñ/°×0BÑ0BÀ1Õ0EÀrÈ1ÓM×PÑPÔW\×WbÑWbÐkl×ksÑksÐPÓtð 	ð !-¨ZÕ 8× >Ñ >Ó @Ðä'1°!·(±(·-±-Ä×'EÒ'EÈ!Ï(É(Ï-É-Ð[`ÔJ`�a—h‘h—m’mÐfkˆä¨¸U×CÖCØ&Õ>×IÑIÈ!ÈQÓOˆEÜ—)’)˜U˜N°Ô3ˆCØ—'‘'“)˜d×4Ñ4Õ4ˆCØ—'‘'“)˜d×4Ñ4Õ4ˆC÷	 Dð �v‰v˜AŸG™GˆvÓ$ c§f¡f°1·7±7 fÓ&;Ð;Ð;÷ D×Cús   Ã(A4FÆF,	)rƒ   rv   r…   r   r‡   r€   rz   ©N)r=   r>   r?   r@   r   rO   Ústaticmethodr‚   r3   Úno_gradr   rg   rC   rD   rn   ro   s   @@r8   rq   rq   x   su   ù‡ € Ù�X vÔ.÷Oô Oó /ðOð  Ù�X vÔ.÷5ñ 5ó /ó ð5ð$ ‡]‚]ƒ_Øñ<ó ó ÷<ð <r;   rq   c                   ó8   a a€ ] tR t^±t oV 3R ltR tRtVtV ;t# )ÚCsmMLPc                ó  <€ \         SV `  4        Wn        VP                  V n        VP                  V n        \
        P                  ! V P                  V P                  VP                  R 7      V n        \
        P                  ! V P                  V P                  VP                  R 7      V n	        \
        P                  ! V P                  V P                  VP                  R 7      V n
        \        VP                  ,          V n        R# )©ÚbiasN)rN   rO   rv   rV   Úintermediate_sizerP   ÚLinearÚmlp_biasÚ	gate_projÚup_projÚ	down_projr   Ú
hidden_actÚact_fn©rU   rv   rW   s   &&€r8   rO   ÚCsmMLP.__init__²   s¶   ø€ Ü‰ÑÔØŒØ!×-Ñ-ˆÔØ!'×!9Ñ!9ˆÔÜŸš 4×#3Ñ#3°T×5KÑ5KÐRX×RaÑRaÔbˆŒÜ—y’y ×!1Ñ!1°4×3IÑ3IÐPV×P_ÑP_Ô`ˆŒÜŸš 4×#9Ñ#9¸4×;KÑ;KÐRX×RaÑRaÔbˆŒÜ˜V×.Ñ.Õ/ˆŽr;   c                ó�   € V P                  V P                  V P                  V4      4      V P                  V4      ,          4      pV# r­   )rº   r¼   r¸   r¹   )rU   r¦   rº   s   && r8   rg   ÚCsmMLP.forward¼   s6   € Ø—N‘N 4§;¡;¨t¯~©~¸aÓ/@Ó#AÀDÇLÁLÐQRÃOÕ#SÓTˆ	ØÐr;   )r¼   rv   rº   r¸   rV   rµ   r¹   ©	r=   r>   r?   r@   rO   rg   rC   rD   rn   ro   s   @@r8   r±   r±   ±   s   ù‡ € õ0÷ò r;   r±   c                óÌ   € V RRV P                   R,          ^,          13,          pV RV P                   R,          ^,          R13,          p\        P                  ! V) V3RR7      # )z*Rotates half the hidden dims of the input..Nrœ   r^   )rj   r3   r£   )r¦   Úx1Úx2s   &  r8   Úrotate_halfrÅ   Á   sZ   € à	
ˆ3Ð"�!—'‘'˜"•+ Õ"Ð"Ð"Õ	#€BØ	
ˆ3�—‘˜•˜qÕ Ñ"Ð"Õ	#€BÜ�9Š9�r�c˜2�Y BÔ'Ð'r;   Úrotary_pos_embc                óÈ   € VP                  V4      pVP                  V4      pW,          \        V 4      V,          ,           pW,          \        V4      V,          ,           pWV3# )aI  Applies Rotary Position Embedding to the query and key tensors.

Args:
    q (`torch.Tensor`): The query tensor.
    k (`torch.Tensor`): The key tensor.
    cos (`torch.Tensor`): The cosine part of the rotary embedding.
    sin (`torch.Tensor`): The sine part of the rotary embedding.
    unsqueeze_dim (`int`, *optional*, defaults to 1):
        The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
        sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
        that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
        k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
        cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
        the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
    `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
)Ú	unsqueezerÅ   )ÚqÚkr¤   r¥   Úunsqueeze_dimÚq_embedÚk_embeds   &&&&&  r8   Úapply_rotary_pos_embrÎ   È   sS   € ð& �-‰-˜Ó
&€CØ
�-‰-˜Ó
&€CØ�wœ; q›>¨CÕ/Õ0€GØ�wœ; q›>¨CÕ/Õ0€GØÐÐr;   c                ód   € V ^8„  d   QhR\         P                  R\        R\         P                  /# )r'   r+   Ún_reprJ   )r3   r[   Úint)r6   s   "r8   r9   r9   â   s.   € ÷ 	Uñ 	UœUŸ\™\ð 	U´#ð 	U¼%¿,¹,ñ 	Ur;   c                ó˜   € V P                   w  r#rEV^8X  d   V # V R,          P                  W#WV4      p V P                  W#V,          WE4      # )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)
)r�   r�   Nr�   r�   )rj   rž   Úreshape)r+   rÐ   ÚbatchÚnum_key_value_headsÚslenr�   s   &&    r8   Ú	repeat_kvr×   â   sU   € ð
 2?×1DÑ1DÑ.€E Ø�„zØÐØ!Ð"2Õ3×:Ñ:¸5ÐW\ÐdlÓm€MØ× Ñ  ¸eÕ(CÀTÓTÐTr;   c                ó   € V ^8„  d   QhR\         P                  R\        P                  R\        P                  R\        P                  R\        P                  R,          R\        R\        R	\
        \        ,          /# )
r'   ÚmoduleÚqueryÚkeyÚvalueÚattention_maskNÚscalingÚdropoutr’   )rP   ÚModuler3   r[   rK   r   r   )r6   s   "r8   r9   r9   î   sw   € ÷ %ñ %Ü�I‰Ið%ä�<‰<ð%ô 
�‰ð%ô �<‰<ð	%ô
 —L‘L 4Õ'ð%ô ð%ô ð%ô Ô'Õ(ñ%r;   c                 ó"  € \        W P                  4      p\        W0P                  4      p	\        P                  ! WP	                  ^^4      4      V,          p
Ve	   W¤,           p
\
        P                  P                  V
R\        P                  R7      P                  VP                  4      p
\
        P                  P                  W¦V P                  R7      p
\        P                  ! W©4      pVP	                  ^^4      P                  4       pWº3# )r'   )r”   r_   )ÚpÚtrainingr^   )r×   Únum_key_value_groupsr3   Úmatmulr¢   rP   Ú
functionalÚsoftmaxra   r`   r_   rß   rã   Ú
contiguous)rÙ   rÚ   rÛ   rÜ   rÝ   rÞ   rß   r’   Ú
key_statesÚvalue_statesÚattn_weightsÚattn_outputs   &&&&&&&,    r8   Úeager_attention_forwardrí   î   sÒ   € ô ˜3× ;Ñ ;Ó<€JÜ˜U×$?Ñ$?Ó@€Lä—<’< ×';Ñ';¸A¸qÓ'AÓBÀWÕL€LØÒ!Ø#Õ4ˆä—=‘=×(Ñ(¨¸2ÄUÇ]Á]Ð(ÓS×VÑVÐW\×WbÑWbÓc€LÜ—=‘=×(Ñ(¨È6Ï?É?Ð(Ó[€LÜ—,’,˜|Ó:€KØ×'Ñ'¨¨1Ó-×8Ñ8Ó:€KàÐ$Ð$r;   c                   óX   a a€ ] tR tRt oRtV3R lV 3R lltRV3R lR lltRtVtV ;t	# )	ÚCsmAttentioni  z=Multi-headed attention from 'Attention Is All You Need' paperc                ó&   <€ V ^8„  d   QhRS[ RS[/# ©r'   rv   Ú	layer_idx©r    rÑ   )r6   r7   s   "€r8   r9   ÚCsmAttention.__annotate__  s   ø€ ÷ 
ñ 
™yð 
±Sñ 
r;   c                ó¢  <€ \         SV `  4        Wn        W n        \	        VR VP
                  VP                  ,          4      V n        VP                  VP                  ,          V n	        V P                  R,          V n
        VP                  V n        RV n        \        P                  ! VP
                  VP                  V P                  ,          VP                  R7      V n        \        P                  ! VP
                  VP                  V P                  ,          VP                  R7      V n        \        P                  ! VP
                  VP                  V P                  ,          VP                  R7      V n        \        P                  ! VP                  V P                  ,          VP
                  VP                  R7      V n        R# )r�   Tr³   Ng      à¿)rN   rO   rv   rò   r�   rV   r�   r�   rÕ   rä   rÞ   Úattention_dropoutÚ	is_causalrP   r¶   Úattention_biasÚq_projÚk_projÚv_projÚo_proj©rU   rv   rò   rW   s   &&&€r8   rO   ÚCsmAttention.__init__  sI  ø€ Ü‰ÑÔØŒØ"ŒÜ ¨
°F×4FÑ4FÈ&×JdÑJdÕ4dÓeˆŒØ$*×$>Ñ$>À&×B\ÑB\Õ$\ˆÔ!Ø—}‘} dÕ*ˆŒØ!'×!9Ñ!9ˆÔØˆŒä—i’iØ×Ñ × :Ñ :¸T¿]¹]Õ JÐQW×QfÑQfô
ˆŒô —i’iØ×Ñ × :Ñ :¸T¿]¹]Õ JÐQW×QfÑQfô
ˆŒô —i’iØ×Ñ × :Ñ :¸T¿]¹]Õ JÐQW×QfÑQfô
ˆŒô —i’iØ×&Ñ&¨¯©Õ6¸×8JÑ8JÐQW×QfÑQfô
ˆŽr;   c                ó  <€ V ^8„  d   QhRS[ P                  RS[S[ P                  S[ P                  3,          R,          RS[ P                  R,          RS[R,          RS[S[,          RS[S[ P                  S[ P                  3,          /# )r'   r+   Úposition_embeddingsNrÝ   r*   r’   rJ   )r3   r[   r5   r   r   r   )r6   r7   s   "€r8   r9   rô   "  s‡   ø€ ÷ &)ñ &)á—|‘|ð&)ñ #¡5§<¡<±·±Ð#=Õ>ÀÕEð&)ñ Ÿ™ tÕ+ð	&)ñ
  �ð&)ñ Ñ+Õ,ð&)ñ 
‰u�|‰|™UŸ\™\Ð)Õ	*ñ&)r;   c                ó  € VP                   R R p. VORNV P                  N5pV P                  V4      P                  V4      P	                  ^^4      pV P                  V4      P                  V4      P	                  ^^4      p	V P                  V4      P                  V4      P	                  ^^4      p
Vw  r¼\        W‰W¼4      w  r‰Ve   VP                  WšV P                  4      w  rš\        P                  ! V P                  P                  \        4      pT! V VV	V
V3RV P                  '       g   RMV P                   RV P"                  /VB w  rïVP$                  ! . VORN5!  P'                  4       pV P)                  V4      pWï3# )Nrß   ç        rÞ   r^   )rj   r�   rù   Úviewr¢   rú   rû   rÎ   Úupdaterò   r   Úget_interfacerv   Ú_attn_implementationrí   rã   rö   rÞ   rÓ   rè   rü   )rU   r+   r   rÝ   r*   r’   Úinput_shapeÚhidden_shapeÚquery_statesré   rê   r¤   r¥   Úattention_interfacerì   rë   s   &&&&&,          r8   rg   ÚCsmAttention.forward"  s‡  € ð $×)Ñ)¨#¨2Ð.ˆØ8˜Ð8 bÐ8¨$¯-©-Ñ8ˆà—{‘{ =Ó1×6Ñ6°|ÓD×NÑNÈqÐRSÓTˆØ—[‘[ Ó/×4Ñ4°\ÓB×LÑLÈQÐPQÓRˆ
Ø—{‘{ =Ó1×6Ñ6°|ÓD×NÑNÈqÐRSÓTˆà&‰ˆÜ#7¸ÐRUÓ#[Ñ ˆàÒ&Ø'6×'=Ñ'=¸jÐX\×XfÑXfÓ'gÑ$ˆJä(?×(MÒ(MØ�K‰K×,Ñ,Ô.Eó)
Ðñ %8ØØØØØñ	%
ð  $Ÿ}Ÿ}˜}‘C°$×2HÑ2Hð	%
ð —L‘Lð	%
ð ñ	%
Ñ!ˆð "×)Ò)Ð;¨;Ð;¸Ó;×FÑFÓHˆØ—k‘k +Ó.ˆØÐ(Ð(r;   )rö   rv   r�   r÷   rú   rò   rä   rü   rù   rÞ   rû   )NNN)
r=   r>   r?   r@   rA   rO   rg   rC   rD   rn   ro   s   @@r8   rï   rï     s!   ù‡ € áG÷
ó 
÷.&)÷ &)ò &)r;   rï   c                   óT   a a€ ] tR tRt oV3R lV 3R lltRV3R lR lltRtVtV ;t# )ÚCsmDecoderLayeriK  c                ó&   <€ V ^8„  d   QhRS[ RS[/# rñ   ró   )r6   r7   s   "€r8   r9   ÚCsmDecoderLayer.__annotate__L  s   ø€ ÷ `ñ `™yð `±Sñ `r;   c                ó   <€ \         SV `  4        VP                  V n        \        WR 7      V n        \        V4      V n        \        VP                  VP                  R7      V n	        \        VP                  VP                  R7      V n
        R# ))rv   rò   ©rI   N)rN   rO   rV   rï   Ú	self_attnr±   ÚmlprG   Úrms_norm_epsÚinput_layernormÚpost_attention_layernormrý   s   &&&€r8   rO   ÚCsmDecoderLayer.__init__L  si   ø€ Ü‰ÑÔØ!×-Ñ-ˆÔä%¨VÔIˆŒä˜&“>ˆŒÜ)¨&×*<Ñ*<À&×BUÑBUÔVˆÔÜ(2°6×3EÑ3EÈ6×K^ÑK^Ô(_ˆÖ%r;   c                ó.  <€ V ^8„  d   QhRS[ P                  RS[ P                  R,          RS[ P                  R,          RS[R,          RS[R,          RS[S[ P                  S[ P                  3,          R,          RS[S[,          R	S[ P                  /# )
r'   r+   rÝ   Nr§   r*   Ú	use_cacher   r’   rJ   )r3   r[   Ú
LongTensorr   Úboolr5   r   r   )r6   r7   s   "€r8   r9   r  V  sœ   ø€ ÷ ñ á—|‘|ðñ Ÿ™ tÕ+ðñ ×&Ñ&¨Õ-ð	ñ
  �ðñ ˜$•;ðñ #¡5§<¡<±·±Ð#=Õ>ÀÕEðñ Ñ+Õ,ðñ 
�‰ñr;   c                óÔ   € TpV P                  V4      pV P                  ! RR VRVRVRVRVRV/VB w  rW�,           pTpV P                  V4      pV P                  V4      pW�,           pV# )r+   rÝ   r§   r*   r  r   r<   )r  r  r  r  )
rU   r+   rÝ   r§   r*   r  r   r’   ÚresidualÚ_s
   &&&&&&&,  r8   rg   ÚCsmDecoderLayer.forwardV  s©   € ð !ˆØ×,Ñ,¨]Ó;ˆàŸ>š>ñ 
Ø'ð
à)ð
ð &ð
ð ,ð	
ð
  ð
ð !4ð
ð ñ
Ñˆð !Õ0ˆð !ˆØ×5Ñ5°mÓDˆØŸ™ Ó/ˆØ Õ0ˆØÐr;   )rV   r  r  r  r  )NNNFNrÁ   ro   s   @@r8   r  r  K  s    ù‡ € ÷`ó `÷÷ ò r;   r  z[
    The bare Csm Model outputting raw hidden-states without any specific head on top.
    c                   óš   a a€ ] tR tRt oRtRtRtR.tR.tRt	Rt
RtRtR]R]/t]P"                  ! 4       V 3R l4       tV3R	 ltR
tVtV ;t# )ÚCsmPreTrainedModeliv  ÚmodelTr  r*   r+   r,   c                ó  <€ \         SV `  V4       \        V\        4      '       d_   VP                  p\        V^,
          4       F:  p\        P                  ! VP                  RV P                  P                  R7       K<  	  R# \        V\        4      '       dg   \        P                  ! VP                  \        P                  ! V P                  P                  4      V P                  P                   ,          4       R# R# )é   r  )rc   ÚstdN)rN   Ú_init_weightsrŸ   ÚCsmCodebooksHeadÚnum_codebooksÚrangeÚinitÚnormal_rS   rv   Úinitializer_rangeÚCsmBackboneModelEmbeddingsÚcopy_Úaudio_tokens_offsetsr3   r‘   Ú
vocab_size)rU   rÙ   r(  ÚirW   s   &&  €r8   r&  Ú CsmPreTrainedModel._init_weights�  s§   ø€ ä‰Ñ˜fÔ%Ü�fÔ.×/Ò/Ø"×0Ñ0ˆMÜ˜=¨1Õ,Ö-�Ü—’˜VŸ]™]°¸$¿+¹+×:WÑ:W×Xó .ä˜Ô :×;Ò;Ü�JŠJ�v×2Ñ2´E·L²LÀÇÁ×AZÑAZÓ4[Ð^b×^iÑ^i×^tÑ^tÕ4tÖuñ <r;   c                ó&   <€ V ^8„  d   Qh/ S[ ;R&   # ru   rw   )r6   r7   s   "€r8   r9   ÚCsmPreTrainedModel.__annotate__v  s   ø‡ ‚ ñ Ñò r;   r<   )ÚaudioÚtext)r=   r>   r?   r@   Úbase_model_prefixÚinput_modalitiesÚsupports_gradient_checkpointingÚ_no_split_modulesÚ_skip_keys_device_placementÚ_supports_flash_attnÚ_supports_sdpaÚ_can_compile_fullgraphÚ_supports_attention_backendr  rï   Ú_can_record_outputsr3   r¯   r&  rB   rC   rD   rn   ro   s   @@r8   r!  r!  v  sy   ù‡ € ð  ÐØ(ÐØ&*Ð#Ø*Ð+ÐØ#4Ð"5ÐØÐØ€Nð "ÐØ"&Ðà˜Ø�lðÐð
 ‡]‚]ƒ_ôvó ðv÷5 … r;   r!  c                   ór   a a€ ] tR tRt oV 3R lt]]]RV3R lR ll4       4       4       tV3R lt	Rt
VtV ;t# )ÚCsmDepthDecoderModeliš  c           	     ó„  <€ \         SV `  V4       VP                  V n        VP                  V n        \
        P                  ! VP                  VP                  ,          VP                  4      V n	        \
        P                  ! \        VP                  4       Uu. uF  p\        W4      NK  	  up4      V n        \        VP                   VP"                  R 7      V n        \'        VR7      V n        RV n        \
        P,                  ! VP                  VP                   RR7      V n        V P1                  4        R# u upi )r  ©rv   Fr³   N)rN   rO   Úpad_token_idÚpadding_idxr0  rP   Ú	Embeddingr(  Úbackbone_hidden_sizeÚembed_tokensÚ
ModuleListr)  Únum_hidden_layersr  ÚlayersrG   rV   r  Únormrq   Ú
rotary_embÚgradient_checkpointingr¶   Úinputs_embeds_projectorÚ	post_initrý   s   && €r8   rO   ÚCsmDepthDecoderModel.__init__ž  sî   ø€ Ü‰Ñ˜Ô Ø!×.Ñ.ˆÔØ ×+Ñ+ˆŒÜŸLšL¨&×*>Ñ*>À×ARÑARÕ*RÐU[×UpÑUpÓqˆÔÜ—m’mÜAFÀv×G_ÑG_ÔA`ÓaÑA`°IŒ_˜VÖ/ÑA`Ñaó
ˆŒô ˜v×1Ñ1°v×7JÑ7JÔKˆŒ	Ü,°FÔ;ˆŒØ&+ˆÔ#Ü')§y¢y°×1LÑ1LÈf×N`ÑN`ÐglÔ'mˆÔ$ð 	�‰Öùò bs   ÂD=c                ó6  <€ V ^8„  d   QhRS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[R,          RS[ P                  R,          RS[R,          R	S[S[,          R
S[S[	,          /	# )r'   Ú	input_idsNÚbackbone_last_hidden_staterÝ   r§   r*   Úinputs_embedsr  r’   rJ   )
r3   r  r4   r[   r   r  r   r   r5   r   )r6   r7   s   "€r8   r9   Ú!CsmDepthDecoderModel.__annotate__±  s¹   ø€ ÷ N
ñ N
á×#Ñ# dÕ*ðN
ñ %*×$5Ñ$5¸Õ$<ðN
ñ Ÿ™ tÕ+ð	N
ñ
 ×&Ñ&¨Õ-ðN
ñ  �ðN
ñ ×(Ñ(¨4Õ/ðN
ñ ˜$•;ðN
ñ Ñ+Õ,ðN
ñ 
Ñ(Õ	(ñN
r;   c                ó`  € Ve(   \        4       '       g   \        P                  R4       RpVRJ VRJ,          '       d   \        R4      hV'       d   Vf   \	        V P
                  R7      pVe   VP                  4       M^ p	Ve   VP                  ^,          MVP                  ^,          p
Ve   VP                  MVP                  p\        P                  ! W™V
,           VR7      pVf‹   \        P                  ! V^,
          ^ R7      pWÀP                  ,          pV P                  W,           4      pV^ ,          ^ 8H  pVe   W&R&   M-\        4       '       g   V'       d   \        P                  R4       V P                  V4      p\!        V P
                  VVVVR7      pTpVP#                  ^ 4      pV P%                  VVR	7      pV P&                  RV P
                  P(                    F  pV! V3R
VRVRVRVRV/VB pK  	  V P+                  V4      p\-        TV'       d   VR7      # RR7      # )a*  
backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*):
    The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model)
    is provided in the `input_ids` argument.
NzÕCustom `position_ids` were provided but will be ignored. CSM depth decoder automatically determines position_ids and as it requires them to be identical across the batch, the provided position_ids will be ignored.z;You must specify exactly one of input_ids or inputs_embeds.rD  ©rr   )ÚminzvWhen the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference.©rv   rV  rÝ   r*   r§   ©r§   rÝ   r§   r*   r  r   ©Úlast_hidden_stater*   )r�   r   )r   ÚloggerÚwarning_onceÚ
ValueErrorr   rv   Úget_seq_lengthrj   rr   r3   r‘   Úclampr0  rI  ÚwarningrP  r   rÈ   rN  rL  rK  rM  r   )rU   rT  rU  rÝ   r§   r*   rV  r  r’   Úpast_seen_tokensÚinputs_seq_lengthrr   Úcodebook_idxsÚoffsetÚinput_ids_are_first_codebookÚcausal_maskr+   r   Údecoder_layers   &&&&&&&&,          r8   rg   ÚCsmDepthDecoderModel.forward®  s;  € ð& Ò#Ô,D×,FÒ,FÜ×Ñðwôð  ˆLØ˜Ð -°tÐ";×<Ô<ÜÐZÓ[Ð[ç˜Ò0Ü*°$·+±+Ô>ˆOà?NÒ?Z˜?×9Ñ9Ô;Ð`aÐØ6CÒ6O˜M×/Ñ/°Ö2ÐU^×UdÑUdÐefÕUgÐØ)6Ò)B�×%Ò%È	×HXÑHXˆÜ—|’|Ð$4ÐIZÕ6ZÐciÔjˆàÒ Ü!ŸKšK¨°qÕ(8¸aÔ@ˆMØ"§_¡_Õ4ˆFØ ×-Ñ-¨iÕ.@ÓAˆMà+7¸­?¸aÑ+?Ð(Ø)Ò5Ø&@˜dÒ#ä/×1Ò1×6RÜ—N‘Nð Qôð ×4Ñ4°]ÓCˆä(Ø—;‘;Ø'Ø)Ø+Ø%ô
ˆð &ˆð $×-Ñ-¨aÓ0ˆØ"Ÿo™o¨mÈ,˜oÓWÐà!Ÿ[™[Ð)H¨4¯;©;×+HÑ+HÓIˆMÙ)Øñà*ðð *ðð !0ð	ð
 $ðð %8ðð ñŠMñ Jð Ÿ	™	 -Ó0ˆÜ&Ø+ß/8˜Oô
ð 	
à>Bô
ð 	
r;   c                ó&   <€ V ^8„  d   Qh/ S[ ;R&   # ru   )r!   )r6   r7   s   "€r8   r9   rW  š  s   ø‡ ‚ á!Ñ!ò r;   )rI  rO  rP  rL  rM  rF  rN  r0  )NNNNNNN)r=   r>   r?   r@   rO   r   r   r   rg   rB   rC   rD   rn   ro   s   @@r8   rB  rB  š  s=   ù‡ € õð   ØØ÷N
ñ N
ó ó ó  ðN
÷/ … r;   rB  c                   ó<   a a€ ] tR tRt oV 3R ltRR ltRtVtV ;t# )r'  i  c                ó²   <€ \         SV `  4        W n        \        P                  ! \
        P                  ! V P                  ^,
          W4      4      V n        R# )r$  N)rN   rO   r(  rP   rQ   r3   ÚemptyrS   )rU   rV   r(  r0  rW   s   &&&&€r8   rO   ÚCsmCodebooksHead.__init__  s:   ø€ Ü‰ÑÔØ*ÔÜ—l’l¤5§;¢;¨t×/AÑ/AÀAÕ/EÀ{Ó#_Ó`ˆŽr;   c           	     ó>  € V^,
          pV P                   V,          p\        VP                  ^ ,          4       Uu. uF=  p\        P                  P                  VRVR3,          W4,          P                  4      NK?  	  pp\        P                  ! V^R7      pV# u upi )r$  r�   rœ   )	rS   r)  rj   rP   ræ   ÚlinearÚTr3   Ústack)rU   r+   Úcodebook_indicesÚcodebook_weightÚcodebook_idxs   &&&  r8   rg   ÚCsmCodebooksHead.forward  s‘   € à+¨aÕ/ÐØŸ+™+Ð&6Õ7ˆô !& o×&;Ñ&;¸AÕ&>Ô ?ó
á ?�ô �M‰M× Ñ  ¨q°,ÀÐ/AÕ!BÀOÕDa×DcÑDcÖdÙ ?ð 	ð 
ô Ÿš M°qÔ9ˆàÐùò
s   »AB)r(  rS   r­   rÁ   ro   s   @@r8   r'  r'    s   ù‡ € õa÷
ô r;   r'  a$  
    The CsmDepthDecoder Model transformer, with a [`CsmCodebooksHead`] on top,
    which can be seen a position-specific language modeling head, allowing to use a different linear layer for each codebook
    (e.g. position 0 is the first codebook and uses the first codebook head, etc.)
    c                   óŒ   a a€ ] tR tRt oRtRtRtRR/tV 3R lt]	]
RV3R lR ll4       4       tRV3R lV 3R	 llltR
tVtV ;t# )ÚCsmDepthDecoderForCausalLMi  NÚlm_headÚkeep_full_weightc                óè   <€ \         SV `  V4       \        V4      V n        VP                  V n        \        VP                  VP                  VP                  4      V n        V P                  4        R # r­   )
rN   rO   rB  r"  r0  r'  rV   r(  Úcodebooks_headrQ  r½   s   &&€r8   rO   Ú#CsmDepthDecoderForCausalLM.__init__#  sY   ø€ Ü‰Ñ˜Ô Ü)¨&Ó1ˆŒ
Ø ×+Ñ+ˆŒÜ.¨v×/AÑ/AÀ6×CWÑCWÐY_×YjÑYjÓkˆÔð 	�‰Ör;   c                óˆ  <€ V ^8„  d   QhRS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[R,          RS[ P                  R,          RS[ P                  R,          R	S[R,          R
S[S[ P                  ,          RS[S[,          RS[	S[
,          /# )r'   rT  NrU  rÝ   r§   r*   rV  Úlabelsr  Úlogits_to_keepr’   rJ   )r3   r  r4   r[   r   r  rÑ   r   r   r5   r   )r6   r7   s   "€r8   r9   Ú'CsmDepthDecoderForCausalLM.__annotate__.  sã   ø€ ÷ A
ñ A
á×#Ñ# dÕ*ðA
ñ %*×$5Ñ$5¸Õ$<ðA
ñ Ÿ™ tÕ+ð	A
ñ
 ×&Ñ&¨Õ-ðA
ñ  �ðA
ñ ×(Ñ(¨4Õ/ðA
ñ × Ñ  4Õ'ðA
ñ ˜$•;ðA
ñ ™eŸl™lÕ*ðA
ñ Ñ+Õ,ðA
ñ 
Ñ'Õ	'ñA
r;   c
                ó  € Ve   VP                  4       M^ pVe   VP                  ^,          MVP                  ^,          pVe   VP                  MVP                  p\        P                  ! WÍR7      V,           pV P
                  ! RRVRVRVRVRVRVR	V/V
B pV^ ,          p\        V	\        4      '       d#   V	^ 8X  d   \        ^R4      pM\        V	) R4      pMT	pV P                  VR
VR
3,          VV,          4      pVP                  4       pRpVeF   VR,          P                  4       pV P                  ! RRVRRRV P                  P                  RV/V
B p\        VVVP                  VP                   VP"                  R7      # )a®  
backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*):
    The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model)
    is provided in the `input_ids` argument.
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]`.
NrY  rT  rU  rÝ   r§   r*   rV  r  r�   r)   r‚  r0  Úshift_labels)r(   r)   r*   r+   r,   r<   ).ºr$  NN)rb  rj   rr   r3   r‘   r"  rŸ   rÑ   Úslicer  rè   Úloss_functionrv   r0  r   r*   r+   r,   )rU   rT  rU  rÝ   r§   r*   rV  r‚  r  rƒ  r’   re  Úseq_lenrr   rv  Úoutputsr+   Úslice_indicesr)   r(   r†  s   &&&&&&&&&&,          r8   rg   Ú"CsmDepthDecoderForCausalLM.forward,  sÅ  € ð0 @OÒ?Z˜?×9Ñ9Ô;Ð`aÐØ,9Ò,E�-×%Ñ% aÖ(È9Ï?É?Ð[\ÕK]ˆØ)6Ò)B�×%Ò%È	×HXÑHXˆÜ Ÿ<š<¨Ô?ÐBRÕRÐà—*’*ñ 	
Øð	
à'Að	
ð *ð	
ð &ð		
ð
 ,ð	
ð (ð	
ð  ð	
ð ñ	
ˆð   �
ˆä�n¤c×*Ò*Ø Ô"ä % a¨£‘ä % ~ o°tÓ <‘à*ˆMà×$Ñ$ ]°1°mÀQÐ3FÕ%GÐIYÐZgÕIhÓiˆØ×"Ñ"Ó$ˆàˆØÒØ! '�?×5Ñ5Ó7ˆLØ×%Ò%ñ ØðØ&*ðØ7;·{±{×7MÑ7MðØ\hðØlrñˆDô &ØØØ#×3Ñ3Ø!×/Ñ/Ø×)Ñ)ô
ð 	
r;   c                óÀ   <€ V ^8„  d   QhRS[ P                  RS[R,          RS[R,          RS[ P                  R,          RS[ P                  R,          RS[R,          /# )r'   rT  Únext_sequence_lengthNr*   rÝ   rV  Úis_first_iteration)r3   r  rÑ   r   r4   r  )r6   r7   s   "€r8   r9   r„  q  so   ø€ ÷ ñ á×#Ñ#ðñ " D�jðñ  �ð	ñ
 ×(Ñ(¨4Õ/ðñ ×(Ñ(¨4Õ/ðñ ! 4�Kñr;   c                ó‚   <€ \         S	V `  ! WW4V3/ VB pV'       g   VP                  R 4       VP                  R4       V# )rU  r§   )rN   Úprepare_inputs_for_generationÚpop)
rU   rT  r�  r*   rÝ   rV  r�  r’   Úmodel_inputsrW   s
   &&&&&&&, €r8   r’  Ú8CsmDepthDecoderForCausalLM.prepare_inputs_for_generationq  sM   ø€ ô ‘wÒ<Ø¨_Èmñ
Ø_eñ
ˆ÷ "Ø×ÑÐ9Ô:ð 	×Ñ˜Ô(àÐr;   )r  r"  r0  )	NNNNNNNNr   )NNNNF)r=   r>   r?   r@   Ú_tied_weights_keysÚ_tp_planÚ_pp_planÚ
_fsdp_planrO   r   r   rg   r’  rC   rD   rn   ro   s   @@r8   r{  r{    s[   ù‡ € ð ÐØ€HØ€HØÐ/Ð0€Jõð Ø÷A
ñ A
ó ó ðA
÷F÷ õ r;   r{  c                   ó8   a a€ ] tR tRt oV 3R ltR tRtVtV ;t# )r-  iˆ  c                ó>  <€ \         SV `  4        \        P                  ! VP                  VP
                  ,          VP                  4      V n        \        P                  ! \        P                  ! VP                  4      VP
                  ,          R R7      V n        R# )Fr|   N)rN   rO   rP   rG  r(  Úcodebook_sizerV   Úembed_audio_tokensr„   r3   r‘   r/  r½   s   &&€r8   rO   Ú#CsmBackboneModelEmbeddings.__init__‰  sj   ø€ Ü‰ÑÔÜ"$§,¢,°×0DÑ0DÀv×G[ÑG[Õ0[Ð^d×^pÑ^pÓ"qˆÔÜ$&§I¢IÜ�LŠL˜×-Ñ-Ó.°×1EÑ1EÕEÐRWô%
ˆÖ!r;   c                ól   € V P                  WP                  ,           4      pVP                  ^R7      pV# )r'   rœ   )r�  r/  Úsum)rU   rT  rV  s   && r8   rg   Ú"CsmBackboneModelEmbeddings.forward�  s4   € Ø×/Ñ/°	×<UÑ<UÕ0UÓVˆØ%×)Ñ)¨aÐ)Ó0ˆØÐr;   )r/  r�  rÁ   ro   s   @@r8   r-  r-  ˆ  s   ù‡ € õ
÷ò r;   r-  c                   óf   a a€ ] tR tRt oV 3R lt]]]RV3R lR ll4       4       4       tRt	Vt
V ;t# )ÚCsmBackboneModeli–  c           	     ó¾  <€ \         SV `  V4       VP                  V n        VP                  V n        \        V4      V n        \        P                  ! \        VP                  4       Uu. uF  p\        W4      NK  	  up4      V n        \        VP                  VP                  R 7      V n        \#        VR7      V n        RV n        V P)                  4        R# u upi )r  rD  FN)rN   rO   rE  rF  r0  r-  rI  rP   rJ  r)  rK  r  rL  rG   rV   r  rM  rq   rN  rO  rQ  rý   s   && €r8   rO   ÚCsmBackboneModel.__init__˜  s®   ø€ Ü‰Ñ˜Ô Ø!×.Ñ.ˆÔØ ×+Ñ+ˆŒÜ6°vÓ>ˆÔÜ—m’mÜAFÀv×G_ÑG_ÔA`ÓaÑA`°IŒ_˜VÖ/ÑA`Ñaó
ˆŒô ˜v×1Ñ1°v×7JÑ7JÔKˆŒ	Ü,°FÔ;ˆŒØ&+ˆÔ#ð 	�‰Öùò bs   Á*Cc                óþ   <€ V ^8„  d   QhRS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[R,          RS[ P                  R,          RS[R,          RS[S[,          R	S[/# )
r'   rT  NrÝ   r§   r*   rV  r  r’   rJ   )	r3   r  r[   r   r4   r  r   r   r   )r6   r7   s   "€r8   r9   ÚCsmBackboneModel.__annotate__ª  s“   ø€ ÷ >
ñ >
á×#Ñ# dÕ*ð>
ñ Ÿ™ tÕ+ð>
ñ ×&Ñ&¨Õ-ð	>
ñ
  �ð>
ñ ×(Ñ(¨4Õ/ð>
ñ ˜$•;ð>
ñ Ñ+Õ,ð>
ñ 
!ñ>
r;   c                óˆ  € VRJ VRJ,          '       d   \        R4      hVf   V P                  V4      pV'       d   Vf   \        V P                  R7      pVfb   Ve   VP	                  4       M^ p\
        P                  ! VP                  ^,          VP                  R7      V,           pVP                  ^ 4      p\        V P                  VVVVR7      p	Tp
V P                  W£R7      pV P                  RV P                  P                    F  pV! V
3RV	RVR	VR
VRV/VB p
K  	  V P                  V
4      p
\        V
VR7      # )aæ  
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`):
    1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input
    requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens.

    2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users.

    Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
    [`PreTrainedTokenizer.__call__`] for details.

    [What are input IDs?](../glossary#input-ids)
Nz:You must specify exactly one of input_ids or inputs_embedsrD  rY  r[  r\  rÝ   r   r§   r*   r  r]  )ra  rI  r   rv   rb  r3   r‘   rj   rr   rÈ   r   rN  rL  rK  rM  r   )rU   rT  rÝ   r§   r*   rV  r  r’   re  rj  r+   r   rk  s   &&&&&&&,     r8   rg   ÚCsmBackboneModel.forward§  se  € ð2 ˜Ð -°tÐ";×<Ô<ÜÐYÓZÐZàÒ Ø*.×*;Ñ*;¸IÓ*FˆMç˜Ò0Ü*°$·+±+Ô>ˆOàÒØCRÒC^˜×=Ñ=Ô?ÐdeÐÜ Ÿ<š<¨×(;Ñ(;¸AÕ(>À}×G[ÑG[Ô\Ð_oÕoˆLØ'×1Ñ1°!Ó4ˆLä(Ø—;‘;Ø'Ø)Ø+Ø%ô
ˆð &ˆØ"Ÿo™o¨m˜oÓWÐà!Ÿ[™[Ð)H¨4¯;©;×+HÑ+HÓIˆMÙ)Øñà*ðð %8ðð *ð	ð
 !0ðð $ðð ñŠMñ Jð Ÿ	™	 -Ó0ˆÜ&Ø+Ø+ô
ð 	
r;   )rI  rO  rL  rM  rF  rN  r0  )NNNNNN)r=   r>   r?   r@   rO   r   r   r   rg   rC   rD   rn   ro   s   @@r8   r£  r£  –  s8   ù‡ € õð  ØØ÷>
ñ >
ó ó ó  ÷>
ð >
r;   r£  zË
    The Csm model consists of two llama-like auto-regressive transformer models: a backbone model that predicts the first codebook token and a depth decoder that predicts the other codebook tokens.
    c                   óÄ   a a€ ] tR tRt oRR/tV 3R ltR tR t]V 3R l4       t	V 3R lt
RV3R	 lR
 lltRV3R lV 3R lllt]]RV3R lR ll4       4       tRtVtV ;t# )ÚCsmForConditionalGenerationië  z5backbone_model.embed_tokens.embed_audio_tokens.weightz'depth_decoder.model.embed_tokens.weightc                óò  <€ \         SV `  V4       VP                  V n        \        P                  ! VP
                  VP                  R R7      V n        \        P                  ! VP                  VP
                  4      V n	        \        P                  V4      V n        \        P                  VP                  4      V n        \         P"                  ! VP$                  4      V n        V P)                  4        R# )Fr³   N)rN   rO   r0  rP   r¶   rV   r|  rG  Útext_vocab_sizeÚembed_text_tokensr£  Ú_from_configÚbackbone_modelr{  Údepth_decoder_configÚdepth_decoderr   Úfrom_configÚcodec_configÚcodec_modelrQ  r½   s   &&€r8   rO   Ú$CsmForConditionalGeneration.__init__õ  s«   ø€ Ü‰Ñ˜Ô Ø ×+Ñ+ˆŒÜ—y’y ×!3Ñ!3°V×5FÑ5FÈUÔSˆŒÜ!#§¢¨f×.DÑ.DÀf×FXÑFXÓ!YˆÔÜ.×;Ñ;¸FÓCˆÔÜ7×DÑDÀV×E`ÑE`ÓaˆÔÜ$×0Ò0°×1DÑ1DÓEˆÔØ�‰Ör;   c                ó.   € V P                   P                  # r­   ©r°  rI  rk   s   &r8   Úget_input_embeddingsÚ0CsmForConditionalGeneration.get_input_embeddingsÿ  s   € Ø×"Ñ"×/Ñ/Ð/r;   c                ó&   € WP                   n        R # r­   r¸  )rU   rÜ   s   &&r8   Úset_input_embeddingsÚ0CsmForConditionalGeneration.set_input_embeddings  s   € Ø+0×ÑÖ(r;   c                ó  <€ VP                  R R4      '       d   \        S
V `  ! V/ VB w  r4M\        S
V `  ! V/ VB pRp\        V4      p\	        VP
                  4      P                  4        UUu/ uF#  w  rxVP                  V4      '       g   K  WvR VbK%  	  p	pp\	        VP                  P
                  4      P                  RR/V	C4       V	 F  p\        VP
                  WW,           4       K!  	  R V9   d   VX3# V# u uppi )Úoutput_loading_infoFÚdepth_decoder_NÚ_from_model_config)ÚgetrN   Úfrom_pretrainedÚlenÚvarsÚgeneration_configÚitemsÚ
startswithr²  r  Údelattr)ÚclsÚargsr’   r"  Úloading_infoÚprefixÚ
prefix_lenÚattrrÜ   Údepth_decoder_attrsrW   s   &*,       €r8   rÃ  Ú+CsmForConditionalGeneration.from_pretrained  s  ø€ à�:‰:Ð+¨U×3Ò3Ü"'¡'Ò"9¸4Ð"JÀ6Ñ"JÑˆE�<ä‘GÒ+¨TÐ<°VÑ<ˆEð "ˆÜ˜“[ˆ
ô  $ E×$;Ñ$;Ó<×BÑBÔDô
áD‘�Ø�‰˜v×&ô %ˆD�Ð˜uÒ$ÙDð 	ñ 
ô 	ˆU× Ñ ×2Ñ2Ó3×:Ñ:Ð<PÐRWÐ;oÐ[nÐ;oÔpó (ˆDÜ�E×+Ñ+¨V­]Ö;ñ (ð ! FÔ*Ø˜,Ð&Ð&àˆLùó
s   Á0D Â	D c                ó  <€ R pV P                   P                  P                  4       pVP                  RR4       VP	                  4        F"  w  rV\        V P                  W5,           V4       K$  	  \        SV `  ! V/ VB  R# )rÀ  Útransformers_versionN)r²  rÆ  Úto_diff_dictr“  rÇ  ÚsetattrrN   Úsave_pretrained)rU   rË  r’   rÍ  rÐ  rÏ  rÜ   rW   s   &*,    €r8   rÖ  Ú+CsmForConditionalGeneration.save_pretrained   sq   ø€ à!ˆØ"×0Ñ0×BÑB×OÑOÓQÐØ×ÑÐ 6¸Ô=Ø.×4Ñ4Ö6‰KˆDÜ�D×*Ñ*¨F­M¸5ÖAñ 7ô 	‰Ò Ð0¨Ô0r;   c                óâ   <€ V ^8„  d   QhRS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          /# )r'   rT  NÚinput_valuesÚinput_values_cutoffsr‚  rJ   rZ   )r6   r7   s   "€r8   r9   Ú(CsmForConditionalGeneration.__annotate__*  sx   ø€ ÷ PBñ PBá—<‘< $Õ&ðPBñ —l‘l TÕ)ðPBñ $Ÿl™l¨TÕ1ð	PBñ
 —‘˜tÕ#ðPBñ 
�‰˜Õ	ñPBr;   c                ó®  € V P                  V4      pVEe¢   \        P                  P                  VR4      pW3^ 8¬  ,          P	                  4       pWf^ 8„  ,          p\
        P                  ! VP                  4       VP                  R7      P                  \        V4      R4      pWvP                  ^4      8  p\
        P                  ! 4       ;_uu_ 4        . p\        W#4       F»  w  ršWª^ 8¬  ,          p
\        V
P                  ^ ,          ^,
          4       F„  pW«,          pW«^,           ,          pV	RWÍ13,          pV P                   P#                  VP                  ^ 4      4      pVP$                  P'                  ^R4      pVP)                  V^ ,          4       K†  	  K½  	  \        R V 4       4      p\
        P*                  ! V Uu. uF?  p\        P                  P                  V^ ^ ^ VVP                  ^ ,          ,
          34      NKA  	  up4      pV P                   P-                  V4      pRRR4       V P.                  P0                  pVV8H  pV P2                  P5                  X4      pVX,          VV&   \
        P6                  ! ^^V P.                  P8                  3VP                  \
        P:                  R7      V P.                  P<                  ,          pV P2                  P5                  V4      P?                  ^4      pWP.                  P@                  8H  pVPC                  VPE                  4       ^4      VV&   Veu   VP                  R4      PC                  ^^V P.                  P8                  4      pVV,          VV&   VVV&   VR8H  PG                  RR7      pRVV^ ,          V^,          R3&   TpR	VR
V/# u upi   + '       g   i     ELž; i)a8  
Merges the input_ids and input_values to produce a single inputs_embeds tensor:
1 - Infers the codec model on the input_values to retrieve codebook token.
2 - Embeds codebook tokens and places them at the correct positions in the inputs_embeds tensor.
3 - If labels are provided, expands them to match codebook dimensions and position the target codebook tokens in the inputs_embeds tensor.

Args:
    input_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`):
        The input ids to embed.
    input_values (`torch.Tensor` of shape `(batch_size, channels, audio_sequence_length)`):
        The audio input values to embed.
    input_values_cutoffs (`torch.Tensor` of shape `(batch_size, max_num_audio)`):
        The cutoffs of the audio input values relative to its batch index, padded with -1 when no audio.
NrY  .c              3   óF   "  € T F  qP                   ^ ,          x € K  	  R# 5i)r   N)rj   )Ú.0Úels   & r8   Ú	<genexpr>ÚQCsmForConditionalGeneration._merge_input_ids_with_input_values.<locals>.<genexpr>Z  s   é € Ð&OÑ=N°r§x¡x°§{¢{Ó=Nùs   ‚!)rr   r_   T©Úas_tupler‡  rV  r‚  ©r$  r   r^   i›ÿÿÿéœÿÿÿ)$r®  rP   ræ   ÚpadÚdiffr3   r‘   Úmaxrr   rž   rÄ  rÈ   r¯   Úzipr)  rj   rµ  ÚencodeÚaudio_codesr¢   Úappendru  Úget_audio_codes_maskrv   Úaudio_token_idr°  rI  rR   r(  ÚlongÚcodebook_eos_token_idÚsqueezeÚaudio_eos_token_idÚrepeatr   Únonzero)rU   rT  rÙ  rÚ  r‚  rV  Úaudio_lengthsÚinput_values_maskÚaudio_tokens_listÚbatch_input_valuesÚbatch_input_values_cutoffsr1  Ú	start_idxÚend_idxÚaudio_batchÚcodec_outputsÚcodebook_idsÚmax_audio_framesrß  Úbatched_audio_token_idsÚaudio_codes_maskrî  Úaudio_token_maskÚaudio_embedsÚaudio_eos_frame_idsÚaudio_eos_embedsÚaudio_eos_token_maskÚlabels_expandedÚ depth_decoder_ignore_frames_idxss   &&&&&                        r8   Ú"_merge_input_ids_with_input_valuesÚ>CsmForConditionalGeneration._merge_input_ids_with_input_values*  s|  € ð* ×.Ñ.¨yÓ9ˆàÓ#ä#%§=¡=×#4Ñ#4Ð5IÈ6Ó#RÐ Ø0ÈÑ1JÕK×PÑPÓRˆMØ)¸!Ñ*;Õ<ˆMÜ %§¢Ð-A×-EÑ-EÓ-GÐP\×PcÑPcÔ d× kÑ kÜ�MÓ" Bó!Ðð !2×4KÑ4KÈAÓ4NÑ NÐô
 —’—•Ø$&Ð!ÜFIÈ,ÖFmÑBÐ&Ø1KÐjkÑLkÕ1lÐ.Ü"Ð#=×#CÑ#CÀAÕ#FÈÕ#JÖK˜Ø$>Õ$A˜	Ø"<À½UÕ"C˜Ø&8¸¸iÐ>OÐ9OÕ&P˜Ø(,×(8Ñ(8×(?Ñ(?À×@UÑ@UÐVWÓ@XÓ(Y˜Ø'4×'@Ñ'@×'JÑ'JÈ1ÈbÓ'Q˜Ø)×0Ñ0°¸aµÖAó Lñ Gnô $'Ñ&OÑ=NÓ&OÓ#OÐ Ü*/¯+ª+Ù`qÓrÑ`qÐZ\”R—]‘]×&Ñ& r¨A¨q°!Ð5EÈÏÉÐQRÍÕ5SÐ+TÖUÑ`qÑró+Ð'ð $(×#3Ñ#3×#HÑ#HÐIZÓ#[Ð ÷! !ð$ "Ÿ[™[×7Ñ7ˆNØ(¨NÑ:Ðà×.Ñ.×;Ñ;Ð<SÓTˆLØ.:Ð;KÕ.LˆMÐ*Ñ+ô —
’
˜A˜q $§+¡+×";Ñ";Ð<ÀY×EUÑEUÔ]b×]gÑ]gÔhØ—+‘+×3Ñ3õ4ð  ð  $×2Ñ2×?Ñ?Ð@SÓT×\Ñ\Ð]^Ó_Ðà#,·±×0NÑ0NÑ#NÐ Ø2B×2IÑ2IÐJ^×JbÑJbÓJdÐfgÓ2hˆMÐ.Ñ/ð Ò!Ø"(×"2Ñ"2°2Ó"6×"=Ñ"=¸aÀÀDÇKÁK×D]ÑD]Ó"^�Ø4KÐL\Õ4]�Ð 0Ñ1Ø8K�Ð 4Ñ5à4:¸d±N×3KÑ3KÐUYÐ3KÓ3ZÐ0Øpt�Ð @ÀÕ CÐEeÐfgÕEhÐjlÐ lÑmØ(�à °¸&ÐAÐAùò= s÷ !—�ús    ÃC3OÇAN>È"OÎ>OÏO	c                ó¬   <€ V ^8„  d   QhRS[ P                  RS[R,          RS[R,          RS[ P                  R,          RS[ P                  R,          /# )r'   rT  r�  Nr*   rÝ   rV  )r3   r  rÑ   r   r4   )r6   r7   s   "€r8   r9   rÛ  |  sa   ø€ ÷ ñ á×#Ñ#ðñ " D�jðñ  �ð	ñ
 ×(Ñ(¨4Õ/ðñ ×(Ñ(¨4Õ/ñr;   c                óZ  <€ \         S	V `  ! R
R VRVRVRVRV/VB pVeŒ   VP                  ^8X  d{   VP                  R4      fh   V P	                  VVP                  R4      VP                  R4      VP                  R4      R	7      pVP                  RVR,          RVR,          R R/4       V# )rT  r�  r*   rÝ   rV  NrÙ  rÚ  r‚  )rT  rÙ  rÚ  r‚  r<   )rN   r’  ÚndimrÂ  r	  r  )
rU   rT  r�  r*   rÝ   rV  r’   r”  Úmerged_inputsrW   s
   &&&&&&,  €r8   r’  Ú9CsmForConditionalGeneration.prepare_inputs_for_generation|  sâ   ø€ ô ‘wÒ<ñ 
Øð
à!5ð
ð ,ð
ð *ð	
ð
 (ð
ð ñ
ˆð Ò  Y§^¡^°qÔ%8¸\×=MÑ=MÈoÓ=^Ò=fØ ×CÑCØ#Ø#ŸZ™Z¨Ó7Ø%+§Z¡ZÐ0FÓ%GØ—z‘z (Ó+ð	 Dó ˆMð ×ÑØ  -°Õ"@À(ÈMÐZbÕLcÐepÐrvÐwôð Ðr;   c                ó°  <€ V ^8„  d   QhRS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[ P                  R,          RS[R,          RS[ P                  R,          R	S[ P                  R,          R
S[R,          RS[S[ P                  ,          RS[S[,          RS[	S[
,          /# )r'   rT  NrÙ  rÝ   rÚ  r§   r*   rV  r‚  r  rƒ  r’   rJ   )r3   r  r[   r   r4   r  rÑ   r   r   r5   r%   )r6   r7   s   "€r8   r9   rÛ  �  sõ   ø€ ÷ Y
ñ Y
á×#Ñ# dÕ*ðY
ñ —l‘l TÕ)ðY
ñ Ÿ™ tÕ+ð	Y
ñ
 $Ÿl™l¨TÕ1ðY
ñ ×&Ñ&¨Õ-ðY
ñ  �ðY
ñ ×(Ñ(¨4Õ/ðY
ñ × Ñ  4Õ'ðY
ñ ˜$•;ðY
ñ ™eŸl™lÕ*ðY
ñ Ñ+Õ,ðY
ñ 
Ñ"Õ	"ñY
r;   c                ór  € Ve8   VP                   ^8X  d'   V P                  WWH4      pVR,          pVR,          pRpV P                  ! RRVRVRVRVRVRV	/VB pV^ ,          p\        V
\        4      '       d   \        V
) R4      MT
pV P                  VR	VR	3,          4      pRpRpRpRpVEe   VR,          pV P                  ! RR
VRVRV P                  P                  /VB pVR,          R8H  P                  RR7      ( pVV,          RRV P                  P                  ^,
          13,          p\        P                  P                  VR^ R7      pVP                  RR7      pVV^ ,          V^,          ^,
          R	3,          pVV,          pV P                   ! RRVRVRV	RRRV/VB pVP"                  pVV,           p\%        TTTTVP&                  VP(                  VP*                  Ve   VP,                  MRVe   VP&                  MRVe   VP(                  MRVe   VP*                  R7      # RR7      # )a`  
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`):
    1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input
    requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens.

    2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users.

    Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
    [`PreTrainedTokenizer.__call__`] for details.

    [What are input IDs?](../glossary#input-ids)
input_values_cutoffs (`torch.Tensor` of shape `(batch_size, max_num_audio)`, *optional*):
    Specify the end positions of audio segments within each batch entry, relative to the concatenated audio input.
    If a batch entry has fewer segments than the maximum, it is padded with -1. For example, in a batch of 2 sequences
    where the first contains 2 audio segments of length l1, and the second contains 1 audio segment of length l2,
    the input_values_cutoffs would be: [[l1, 2 * l1], [l2, -1]].
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
    Labels for computing the masked language modeling loss. Indices should be in `[config.audio_token_id, -100, -101]`.
    Requires targeted `input_values` to be provided as audio tokens will be inferred from it using the `codec_model`.
    - `config.audio_token_id` indicates an audio frames (considering sequence length elements as frames)
    - `-100` will be ignored in the loss computation
    - `-101` indicates the audio frame will be used only for the backbone model (using the first codebook token as labels)

    Such labels can be prepared using `output_labels=True` when calling [`CsmProcessor`].
logits_to_keep (`int` or `torch.Tensor`, *optional*):
    Kept for compatibility. Does not support another value than:
    1. `0`, which is equivalent to keeping all logits, used in the training regime
    2. `1`, which is equivalent to keeping only the last logit, used in the generation regime

Example:

```python
>>> import torch
>>> from transformers import CsmForConditionalGeneration, AutoProcessor
>>> from datasets import load_dataset, Audio

>>> model_id = "sesame/csm-1b"
>>> torch_device = "cuda" if torch.cuda.is_available() else "cpu"

>>> processor = AutoProcessor.from_pretrained(model_id)

>>> ds = load_dataset("hf-internal-testing/dailytalk-dummy", split="train")
>>> # ensure the audio is 24kHz
>>> ds = ds.cast_column("audio", Audio(sampling_rate=24000))

>>> conversation = []
>>> # prepare a conversation with text and corresponding audio
>>> for text, audio, speaker_id in zip(ds[:4]["text"], ds[:4]["audio"], ds[:4]["speaker_id"]):
...     conversation.append(
...         {
...             "role": f"{speaker_id}",
...             "content": [{"type": "text", "text": text}, {"type": "audio", "path": audio["array"]}],
...         }
...     )

>>> inputs = processor.apply_chat_template(
...     conversation,
...     tokenize=True,
...     return_dict=True,
...     output_labels=True,
... ).to(torch_device)

>>> model = CsmForConditionalGeneration.from_pretrained(model_id, device_map=torch_device)
>>> output = model(**inputs)
>>> output.loss.backward()
```NrV  r‚  rT  rÝ   r§   r*   r  r�   r)   r0  rœ   .)rÜ   Trâ  rU  Úreturn_dict)r(   r2   r-   r)   r*   r+   r,   r.   r/   r0   r1   r<   )r�   r�   r   )r�   r�   r‡  rå  r^   rä  )r  r	  r°  rŸ   rÑ   rˆ  r|  r‰  rv   r0  Úallr(  rP   ræ   ræ  rô  r²  r(   r%   r*   r+   r,   r)   )rU   rT  rÙ  rÝ   rÚ  r§   r*   rV  r‚  r  rƒ  r’   r  Úbackbone_outputsÚbackbone_hidden_statesrŒ  Úbackbone_logitsr(   r2   r-   Údepth_decoder_outputsÚbackbone_labelsÚ
train_maskÚdepth_decoder_input_idsÚ
train_idxsÚbackbone_last_hidden_statesÚdepth_decoder_labelss   &&&&&&&&&&&,               r8   rg   Ú#CsmForConditionalGeneration.forward›  sß  € ðd Ò  Y§^¡^°qÔ%8Ø ×CÑCØÐ)=óˆMð *¨/Õ:ˆMØ" 8Õ,ˆFØˆIà×.Ò.ñ 
Øð
à)ð
ð &ð
ð ,ð	
ð
 (ð
ð  ð
ð ñ
Ðð "2°!Õ!4Ðä8BÀ>ÔSV×8WÒ8Wœ˜~˜o¨tÔ4Ð]kˆØŸ,™,Ð'=¸aÀÐPQÐ>QÕ'RÓSˆàˆØˆØ!ÐØ $ÐØÓà$ W�oˆOØ ×.Ò.ñ Ø&ðØ/>ðØKOÏ;É;×KaÑKaðØekñˆMð " (Õ+¨tÑ3×8Ñ8¸RÐ8Ó@Ð@ˆJØ&,¨ZÕ&8¸Ð>]ÀÇÁ×@YÑ@YÐ\]Õ@]Ð>]Ð9]Õ&^Ð#ä&(§m¡m×&7Ñ&7Ð8OÐQWÐ_`Ð&7Ó&aÐ#à#×+Ñ+°TÐ+Ó:ˆJØ*@ÀÈAÅÐPZÐ[\ÕP]Ð`aÕPaÐcdÐAdÕ*eÐ'Ø#)¨*Õ#5Ð à$(×$6Ò$6ñ %Ø1ð%à+Fð%ð $ð%ð !ð	%ð
 ,ð%ð ñ%Ð!ð "7×!;Ñ!;ÐØ Ð#5Õ5ˆDä ØØ'Ø1Ø"Ø,×<Ñ<Ø*×8Ñ8Ø'×2Ñ2ØAVÒAbÐ!6×!=Ò!=Ðhlà$Ò0ð +@×*OÒ*Oàà$Ò0ð )>×(KÒ(KàØI^ÒIjÐ%:×%EÑ%Eô
ð 	
ð quô
ð 	
r;   )r°  rµ  r²  r®  r|  r0  )NNNN)
NNNNNNNNNr   )r=   r>   r?   r@   r–  rO   r¹  r¼  ÚclassmethodrÃ  rÖ  r	  r’  r   r   rg   rC   rD   rn   ro   s   @@r8   r«  r«  ë  s�   ù‡ € ð 	@ÐAjðÐõò0ò1ð ôó ðõ41÷PBò PB÷dõ ð> Ø÷Y
ñ Y
ó ó ÷Y
ð Y
r;   r«  )r$  )r  )r!  r£  rB  r{  r«  )OÚcollections.abcr   Údataclassesr   r3   Útorch.nnrP   Ú r   r*  Úactivationsr   Úcache_utilsr   r   Ú
generationr	   Úintegrationsr
   r   Úmasking_utilsr   Úmodeling_layersr   Úmodeling_outputsr   r   Úmodeling_rope_utilsr   r   Úmodeling_utilsr   r   Úprocessing_utilsr   Úutilsr   r   r   r   r   Úutils.deprecationr   Úutils.genericr   r   Úutils.import_utilsr   Úutils.output_capturingr   Úautor   Úconfiguration_csmr    r!   Úgeneration_csmr"   Ú
get_loggerr=   r_  r%   rà   rG   rq   r±   rÅ   rÎ   r×   rí   rï   r  r!  rB  r'  r{  r-  r£  r«  Ú__all__r<   r;   r8   Ú<module>r8     s,  ðõ* %Ý !ã Ý å &Ý !ß .Ý )ß LÝ /Ý 9ß Oß Kß FÝ &ß _Õ _Ý 0ß GÝ :Ý 5Ý ß ?Ý .ð 
×	Ò	˜HÓ	%€ñ ðôð
 ô'3˜ó '3ó óð'3ñT ˜YÓ'ôJ�—‘ó Jó (ðJô(6<˜Ÿ™ô 6<ôrˆR�Y‰Yô ò (ñ Ð-Ó.óó /ðõ2	U÷%ñ2 Ð)Ó*ô@)�2—9‘9ó @)ó +ð@)ôF(Ð0ô (ñV ðôð
 ôv˜ó vó óðvð< ôd
Ð-ó d
ó ðd
ôN�r—y‘yô ñ( ðôôhÐ!3°_ó hóðhôV §¡ô ð ôQ
Ð)ó Q
ó ðQ
ñh ðôô
F
Ð"4Ð6Hó F
óð
F
òR
‚r;   