
    ^j0k                        d dl mZ d dlmZ d dlZd dlmc mZ d dlmZ ddl	m
Z
mZ ddlmZ ddlmZmZmZ dd	lmZ dd
lmZmZ ddlmZ ddlmZmZ ddlmZmZ ddlm Z m!Z! ddl"m#Z# ddl$m%Z%m&Z&m'Z' ddl(m)Z)m*Z* ddl+m,Z,m-Z- ddl.m/Z/ ddl0m1Z1  e,       r	d dl2m3Z3m4Z4 nd\  Z3Z4 ed       G d dejj                               Z6 G d dejj                        Z7 G d dejj                        Z8d Z9 ed       d=d!       Z:d"ejv                  d#e<d$ejv                  fd%Z=	 d>d&ejj                  d'ejv                  d(ejv                  d)ejv                  d*ejv                  dz  d+e>d,e>d-e#e%   fd.Z? ee:       G d/ d0ejj                               Z@d1 ZAe3e4fZB eCeB      ZD G d2 d3ejj                        ZE G d4 d5e      ZFe& G d6 d7e!             ZGe& G d8 d9eG             ZHe& G d: d;eGe             ZIg d<ZJy)?    )Callable)OptionalN)nn   )CacheDynamicCache)GenerationMixin)use_kernel_forward_from_hubuse_kernel_func_from_hubuse_kernelized_func)force_accelerate_hooks)create_causal_maskcreate_recurrent_attention_mask)GradientCheckpointingLayer)BaseModelOutputWithPastCausalLMOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringcan_return_tuple)maybe_autocastmerge_with_config_defaults)is_causal_conv1d_availableis_torchdynamo_compiling)capture_outputs   )
Lfm2Config)causal_conv1d_fncausal_conv1d_update)NNRMSNormc                   h     e Zd Zddeddf fdZdej                  dej                  fdZd Z xZ	S )	Lfm2RMSNormepsreturnNc                     t         |           t        j                  t	        j
                  |            | _        || _        y)z:
        Lfm2RMSNorm is equivalent to T5LayerNorm
        N)super__init__r   	Parametertorchonesweightvariance_epsilon)selfhidden_sizer'   	__class__s      q/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/transformers/models/lfm2/modeling_lfm2.pyr+   zLfm2RMSNorm.__init__4   s1     	ll5::k#:; #    hidden_statesc                 "   |j                   }|j                  t        j                        }|j	                  d      j                  dd      }|t        j                  || j                  z         z  }| j                  |j                  |      z  S )N   T)keepdim)	dtypetor-   float32powmeanrsqrtr0   r/   )r1   r6   input_dtypevariances       r4   forwardzLfm2RMSNorm.forward<   sy    #))%((7 $$Q',,R,>%Ht?T?T4T(UU{{]--k:::r5   c                 ^    t        | j                  j                         d| j                   S )Nz, eps=)tupler/   shaper0   )r1   s    r4   
extra_reprzLfm2RMSNorm.extra_reprC   s*    ))*+6$2G2G1HIIr5   )gư>)
__name__
__module____qualname__floatr+   r-   TensorrC   rG   __classcell__r3   s   @r4   r&   r&   2   s7    $ $$ $;U\\ ;ell ;Jr5   r&   c                        e Zd ZU ej                  ed<   ddef fdZe	 	 	 ddedz  de	d   de
dz  ded	ef   fd
       Z ej                         ed               Z xZS )Lfm2RotaryEmbeddinginv_freqNconfigc                    t         |           |j                  | _        |j                  | _        || _        | j
                  j                  d   | _        | j                  }| j                  dk7  rt        | j                     } || j
                  |      \  }| _
        | j                  d|d       | j                  d|j                         d       y )N	rope_typedefaultrQ   F)
persistentoriginal_inv_freq)r*   r+   max_position_embeddingsmax_seq_len_cachedoriginal_max_seq_lenrR   rope_parametersrT   compute_default_rope_parametersr   attention_scalingregister_bufferclone)r1   rR   devicerope_init_fnrQ   r3   s        r4   r+   zLfm2RotaryEmbedding.__init__J   s    "("@"@$*$B$B!44[A!%!E!E>>Y&.t~~>L+7V+L($(ZeD0(..2BuUr5   r`   ztorch.deviceseq_lenr(   ztorch.Tensorc                    | j                   d   }t        | dd      xs | j                  | j                  z  }d}d|t	        j
                  d|dt        j                        j                  |t        j                        |z  z  z  }||fS )	a  
        Computes the inverse frequencies according to the original RoPE implementation
        Args:
            config ([`~transformers.PreTrainedConfig`]):
                The model configuration.
            device (`torch.device`):
                The device to use for initialization of the inverse frequencies.
            seq_len (`int`, *optional*):
                The current sequence length. Unused for this type of RoPE.
        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   r8   r;   )r`   r;   )	r[   getattrr2   num_attention_headsr-   arangeint64r<   rK   )rR   r`   rb   basedimattention_factorrQ   s          r4   r\   z3Lfm2RotaryEmbedding.compute_default_rope_parametersZ   s    & %%l3fj$/c63E3EIcIc3c U\\!S!5;;?BB&X]XcXcBdgjjk
 )))r5   c                 N   | j                   d d d d f   j                         j                  |j                  d   dd      j	                  |j
                        }|d d d d d f   j                         }t        |j
                  j                  t              r/|j
                  j                  dk7  r|j
                  j                  nd}t        |d      5  |j                         |j                         z  j                  dd      }t        j                  ||fd	      }|j                         | j                  z  }|j                         | j                  z  }	d d d        j	                  |j                   
      	j	                  |j                   
      fS # 1 sw Y   AxY w)Nr   r9   r    mpscpuF)device_typeenabledr8   rl   rf   )rQ   rK   expandrF   r<   r`   
isinstancetypestrr   	transposer-   catcosr]   sinr;   )
r1   xposition_idsinv_freq_expandedposition_ids_expandedrq   freqsembrz   r{   s
             r4   rC   zLfm2RotaryEmbedding.forwardx   sR    !MM$4-8>>@GGHZHZ[\H]_acdehhijiqiqr ,QaZ 8 > > @'1!((--'E!((--[`J`ahhmmfkUC 	5&,,.1F1L1L1NNYYZ[]^_E))UEN3C'')d444C'')d444C		5 vvAGGv$cff177f&;;;	5 	5s   BFF$NNNN)rH   rI   rJ   r-   rL   __annotations__r!   r+   staticmethodr   intrE   rK   r\   no_gradr   rC   rM   rN   s   @r4   rP   rP   G   s    llVz V  $(+/"*T!*(* t* 
~u$	%	* *: U]]_<  <r5   rP   c                   *     e Zd Zdef fdZd Z xZS )Lfm2MLPrR   c                    t         |           |j                  }|j                  rat	        d|z  dz        }|j
                  Dt	        |j
                  |z        }|j                  ||j                  z   dz
  |j                  z  z  }t        j                  |j                  |d      | _
        t        j                  |j                  |d      | _        t        j                  ||j                  d      | _        y )Nr8   r   r    Fbias)r*   r+   intermediate_sizeblock_auto_adjust_ff_dimr   block_ffn_dim_multiplierblock_multiple_ofr   Linearr2   w1w3w2)r1   rR   r   r3   s      r4   r+   zLfm2MLP.__init__   s    "44** #A(9$9A$= >..:$'(G(GJ[([$\!$*$<$<&)A)AAAE&JbJbb%! ))F..0AN))F..0AN))-v/A/ANr5   c                     | j                  t        j                  | j                  |            | j	                  |      z        S r   )r   Fsilur   r   )r1   r|   s     r4   rC   zLfm2MLP.forward   s/    wwqvvdggaj)DGGAJ677r5   )rH   rI   rJ   r!   r+   rC   rM   rN   s   @r4   r   r      s    Oz O8r5   r   c                     | dd| j                   d   dz  f   }| d| j                   d   dz  df   }t        j                  | |fd      S )z*Rotates half the hidden dims of the input..Nr9   r8   rs   )rF   r-   ry   )r|   x1x2s      r4   rotate_halfr      sZ    	
3"!''"+"""	#B	
3q ""	#B99rc2YB''r5   rotary_pos_embc                     |j                  |      }|j                  |      }| |z  t        |       |z  z   }||z  t        |      |z  z   }||fS )a  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krz   r{   unsqueeze_dimq_embedk_embeds          r4   apply_rotary_pos_embr      sY    & --
&C
--
&C3w;q>C/0G3w;q>C/0GGr5   r6   n_repr(   c                     | j                   \  }}}}|dk(  r| S | dddddddddf   j                  |||||      } | j                  |||z  ||      S )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    N)rF   rt   reshape)r6   r   batchnum_key_value_headsslenre   s         r4   	repeat_kvr      so    
 2?1D1D.Ehz!!Qa"23::5BUW\^bdlmM  (;e(CT8TTr5   modulequerykeyvalueattention_maskscalingdropoutkwargsc                    t        || j                        }t        || j                        }	t        j                  ||j	                  dd            |z  }
||
|z   }
t
        j                  j                  |
dt        j                        j                  |j                        }
t
        j                  j                  |
|| j                        }
t        j                  |
|	      }|j	                  dd      j                         }||
fS )Nr8   r   r9   )rl   r;   )ptrainingr    )r   num_key_value_groupsr-   matmulrx   r   
functionalsoftmaxr=   r<   r;   r   r   
contiguous)r   r   r   r   r   r   r   r   
key_statesvalue_statesattn_weightsattn_outputs               r4   eager_attention_forwardr      s     3 ; ;<JUF$?$?@L<<z';';Aq'ABWLL!#n4==((2U]](SVVW\WbWbcL==((6??([L,,|\:K''1-88:K$$r5   c                        e Zd ZdZdedef fdZ	 ddej                  de	ej                  ej                  f   dej                  dz  d	e
dz  d
e	ej                  ej                  dz  f   f
dZ xZS )Lfm2Attentionz=Multi-headed attention from 'Attention Is All You Need' paperrR   	layer_idxc                    t         |           || _        || _        t	        |d|j
                  |j                  z        | _        |j                  |j                  z  | _	        | j                  dz  | _
        d| _        t        j                  |j
                  |j                  | j                  z  d      | _        t        j                  |j
                  |j                  | j                  z  d      | _        t        j                  |j
                  |j                  | j                  z  d      | _        t        j                  |j                  | j                  z  |j
                  d      | _        t%        | j                  |j&                        | _        t%        | j                  |j&                        | _        y )Nre   g      TFr   r'   )r*   r+   rR   r   rg   r2   rh   re   r   r   r   	is_causalr   r   q_projk_projv_projout_projr&   norm_epsq_layernormk_layernormr1   rR   r   r3   s      r4   r+   zLfm2Attention.__init__   sL   "
F4F4F&JdJd4de$*$>$>&B\B\$\!}}d*ii 2 2F4N4NQUQ^Q^4^ejkii 2 2F4N4NQUQ^Q^4^ejkii 2 2F4N4NQUQ^Q^4^ejk		&"<"<t}}"LfN`N`glm&t}}&//J&t}}&//Jr5   Nr6   position_embeddingsr   past_key_valuesr(   c                 
   |j                   d d }g |d| j                  }| j                   | j                  |      j                  |       j                  dd      }| j                   | j                  |      j                  |       j                  dd      }	 | j                  |      j                  | j                  dd      }
|\  }}t        ||	||      \  }}	| |j                  |	|
| j                        \  }	}
t        j                  | j                  j                  t               } || ||	|
|fd| j"                  d|\  }} |j$                  g |d j'                         }| j)                  |      }||fS )Nr9   r    r8           )r   r   )rF   re   r   r   viewrx   r   r   r   r   updater   r   get_interfacerR   _attn_implementationr   r   r   r   r   )r1   r6   r   r   r   r   input_shapehidden_shapequery_statesr   r   rz   r{   attention_interfacer   r   outputs                    r4   rC   zLfm2Attention.forward   s    $))#2.88b8$--8''(GM(B(G(G(VWaabcefg%%&Edkk-&@&E&E|&TU__`acde
6t{{=166EOOPQSTU&S#7jRUWZ#[ j&'6'='=j,X\XfXf'g$J(?(M(MKK,,.E)
 %8	%
 LL	%
 	%
!\ *k));;;;FFH{+|##r5   r   )rH   rI   rJ   __doc__r!   r   r+   r-   rL   rE   r   rC   rM   rN   s   @r4   r   r      s    GKz Kc K( )-%$||%$ #5<<#=>%$ t+	%$
 %$ 
u||U\\D00	1%$r5   r   c                     |N|j                   d   dkD  r<|j                   d   dkD  r*| j                  }| |dddddf   z  j                  |      } | S )zm
    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
    Nr    r   )rF   r;   r<   )r6   r   r;   s      r4   apply_mask_to_padding_statesr     sa    
 !n&:&:1&=&AnFZFZ[\F]`aFa##&1d
)CCGGNr5   c                       e Zd Zdedef fdZ	 	 	 ddej                  dedz  dej                  dz  dej                  dz  fd	Z
	 	 	 ddej                  dedz  dej                  dz  dej                  dz  fd
Z ed      	 	 	 ddej                  dedz  dej                  dz  dej                  dz  fd       Z xZS )Lfm2ShortConvrR   r   c           	      Z   t         |           || _        || _        |j                  | _        |j                  | _        t        j                  |j                  |j                  | j
                  |j                  | j                  | j
                  dz
        | _        t        j                  |j                  d|j                  z  | j                        | _        t        j                  |j                  |j                  | j                        | _        |j                  |   | _        y )Nr    )in_channelsout_channelskernel_sizegroupsr   paddingr   r   )r*   r+   rR   r   conv_L_cacheL_cache	conv_biasr   r   Conv1dr2   convr   in_projr   layer_types
layer_typer   s      r4   r+   zLfm2ShortConv.__init__.  s    
 	"**$$	II**++%%LL1$
	 yy!3!3Q9K9K5KRVR[R[\		&"4"4f6H6HtyyY ,,Y7r5   Nr|   r   r   seq_idxc                    t        ||      }| j                  |      j                  dd      }|j                  dd      \  }}}||z  }| j                  j
                  j                  | j                  j
                  j                  d      | j                  j
                  j                  d            }	||j                  | j                        rht        |j                  d      |j                  | j                     j                  d   |	| j                  j                  d       }
|
j                  d      }
n|jt         j"                  j%                  || j&                  |j(                  d   z
  df      }|j+                  || j                        d| j&                   d f   }t-        ||	| j                  j                  d |      }
||
z  }| j/                  |j                  dd      j1                               }|S )	Nr9   r   rs   r   r8   .)
activationr   )r   r   rx   chunkr   r/   r   sizehas_previous_stater   r#   squeezelayersconv_statesr   r   r   r   padr   rF   update_conv_stater"   r   r   )r1   r|   r   r   r   BCxBCBxconv_weightsconv_out
conv_stateys                r4   cuda_kernels_forwardz"Lfm2ShortConv.cuda_kernels_forwardF  s    )N;ll1o''B/))A2)&1aUyy'',,TYY-=-=-B-B1-EtyyGWGWG\G\]^G_`&?+M+Mdnn+]+

2&&t~~6BB1E		H  ))"-H*]]..rDLL288B<4OQR3ST
,>>z4>>Z[^aeamam`m`o[op
 (L$))..UYcjkHLMM!++b"-88:;r5   c           
         |j                   d   }t        ||      }| j                  |      j                  dd      }|j	                  dd      \  }}}||z  }	||j                  | j                        r|j                  |	| j                        d| j                   d f   }
t        j                  |
j                  |	j                        | j                  j                  d d dd d f   z  d      }| j                  r|| j                  j                  z  }|j!                  d      }n|w|j                   d   dk(  rd|jt"        j$                  j'                  |	| j                  |	j                   d   z
  df      }
|j                  |
| j                        d| j                   d f   }
|d   }|dd  |d d k7  j)                  d	      d   dz   }t        j*                  |j-                  d      ||j/                  d
|j1                               g      j3                         }g }t5        t7        |      dz
        D ]K  }||   ||dz      }}||kD  s|j9                  | j                  |	d d d d ||f         dd ||z
  f          M t        j*                  |d      }n|jt"        j$                  j'                  |	| j                  |	j                   d   z
  df      }
|j                  |
| j                        d| j                   d f   }
| j                  |	      dd |f   }||z  }|j                  dd      j;                         }| j=                  |      }|S )Nr    r9   r   r   rs   .r   T)as_tupler    )rF   r   r   rx   r   r   r   r   r   r-   sumr<   r`   r   r/   r   r   r   r   r   nonzerory   	new_zerosnew_fullnumeltolistrangelenappendr   r   )r1   r|   r   r   r   seqlenr   r   r   r   r  r   sichangeboundspartsiser  s                       r4   slow_forwardzLfm2ShortConv.slow_forwardi  s    (N;ll1o''B/))A2)&1aU&?+M+Mdnn+](::2t~~NsUYUaUaTaTcOcdJyyryy!9DII<L<LQPQSTW<U!U[]^HyyDIINN*))"-H QWWQZ1_*]]..rDLL288B<4OQR3ST
,>>z4>>Z[^aeamam`m`o[op
Bf3B'00$0?BQFFYY 0 0 3VV__TSUS[S[S]=^_`ggiFE3v;?+ Iay&Q-1q5LL2aAaCi=!9#wQw,!GHI yyB/H*]]..rDLL288B<4OQR3ST
,>>z4>>Z[^aeamam`m`o[op
yy}S'6'\2HLKKB**,MM!r5   r   r6   c                     t         r7d|j                  j                  v rt               s| j	                  ||||      S | j                  ||||      S )Ncuda)r   )is_fast_path_availabler`   rv   r   r  r  )r1   r6   r   r   r   s        r4   rC   zLfm2ShortConv.forward  sV     "f0D0D0I0I&IRjRl,,]O^el,mm  Y` aar5   r   )rH   rI   rJ   r!   r   r+   r-   rL   r   	IntTensorr  r  r   rC   rM   rN   s   @r4   r   r   -  s&   88 86 )-.2*.!<<! ! t+	!
 4'!L )-.2*..<<. . t+	.
 4'.` F# )-.2*.	b||	b 	b t+		b
 4'	b $	br5   r   c                        e Zd Zdedef fdZ	 	 	 	 ddej                  deej                  ej                  f   dz  dej                  dz  dej                  dz  d	e
dz  d
ej                  fdZ xZS )Lfm2DecoderLayerrR   r   c                 f   t         |           |j                  |   dk(  | _        | j                  rt	        ||      | _        nt        ||      | _        t        |      | _	        t        |j                  |j                        | _        t        |j                  |j                        | _        y )Nfull_attentionr   )r*   r+   r   is_attention_layerr   	self_attnr   r   r   feed_forwardr&   r2   r   operator_normffn_normr   s      r4   r+   zLfm2DecoderLayer.__init__  s    "("4"4Y"?CS"S""*69=DN%fi8DI#FO(););Q#F$6$6FOOLr5   Nr6   r   r   r}   r   r(   c           	      .   |}| j                   r+ | j                  d| j                  |      ||||d|\  }}n3| j                  | j                  |      |||j	                  d            }||z   }|| j                  | j                  |            z   }|S )N)r6   r   r   r}   r   r   )r6   r   r   r    )r!  r"  r$  r   getr#  r%  )	r1   r6   r   r   r}   r   r   residual_s	            r4   rC   zLfm2DecoderLayer.forward  s     !""-t~~  "00?$7-) /   M1 !II"00? /-

9-	 & M &0%(9(9$--:V(WWr5   )NNNN)rH   rI   rJ   r!   r   r+   r-   rL   rE   
LongTensorr   rC   rM   rN   s   @r4   r  r    s    
Mz 
Mc 
M IM.204(,|| #5<<#=>E t+	
 &&-  
r5   r  c                   N    e Zd ZU eed<   dZdZdgZdgZdZ	dZ
dZdZdZeedZdZy)Lfm2PreTrainedModelrR   modelTr  r   )r6   
attentionsN)rH   rI   rJ   r!   r   base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_supports_flex_attn_can_compile_fullgraph_supports_attention_backendr  r   _can_record_outputs_is_statefulr'  r5   r4   r-  r-    sX    &*#+,#4"5N!"&)# Lr5   r-  c                        e Zd Zdef fdZeee	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dedz  dej                  dz  d	edz  d
ee   defd                     Z xZS )	Lfm2ModelrR   c           	         t         |   |       |j                  | _        |j                  | _        t        j                  |j                  |j                  | j                        | _        t        j                  t        |j                        D cg c]  }t        ||       c}      | _        t        |      | _        d| _        t#        |j                  |j$                        | _        | j)                          y c c}w )NrR   Fr   )r*   r+   pad_token_idpadding_idx
vocab_sizer   	Embeddingr2   embed_tokens
ModuleListr  num_hidden_layersr  r   rP   
rotary_embgradient_checkpointingr&   r   embedding_norm	post_initr   s      r4   r+   zLfm2Model.__init__  s     !.. ++LL):):F<N<NPTP`P`ammBGH`H`BabYfi0b
 .V<&+#)&*<*<&//R 	 cs   DN	input_idsr   r}   r   inputs_embeds	use_cacher   r(   c           	         |d u |d uz  rt        d      || j                  |      }|r|t        | j                        }|V||j	                         nd}t        j                  |j                  d   |j                        |z   }|j                  d      }t        |x}	t              s)| j                  ||||d}
t        di |
t        di |
d}	|}| j                  ||      }t        | j                   d | j                  j"                         D ]-  \  }} ||f|	| j                  j$                  |      |||d	|}/ | j'                  |      }t)        ||
      S )Nz:You must specify exactly one of input_ids or inputs_embedsr>  r   r    )r`   )rR   rK  r   r   r}   )r   r   )r}   )r   r   r}   r   )last_hidden_stater   r'  )
ValueErrorrC  r   rR   get_seq_lengthr-   ri   rF   r`   r   ru   dictr   r   rF  	enumerater   rE  r   rH  r   )r1   rJ  r   r}   r   rK  rL  r   past_seen_tokenscausal_mask_mappingmask_kwargsr6   r   r  decoder_layers                  r4   rC   zLfm2Model.forward  s    -t";<YZZ  --i8M0*$++>OCRC^==?de <<(;(;A(>}G[G[\_ooL'11!4L?-F++!."0#2 ,K #5"C{"C7F+F#
 &"oom,oW !*$++6U8U8U*V W 	A})24;;3J3J13MN$7) / M	 ++M:&++
 	
r5   )NNNNNN)rH   rI   rJ   r!   r+   r   r   r   r-   r+  rL   r   FloatTensorboolr   r   r   rC   rM   rN   s   @r4   r<  r<    s    z     .2.204(,26!%8
##d*8
 t+8
 &&-	8

 8
 ((4/8
 $;8
 +,8
 
!8
    8
r5   r<  c                   B    e Zd ZddiZddiZddgdgfiZ fdZee	 	 	 	 	 	 	 	 dd	e	j                  dz  d
e	j                  dz  de	j                  dz  dedz  de	j                  dz  de	j                  dz  dedz  dee	j                  z  dee   defd              Z xZS )Lfm2ForCausalLMzlm_head.weightzmodel.embed_tokens.weightlm_headcolwise_gather_outputr6   logitsc                     t         |   |       t        |      | _        |j                  | _        t        j                  |j                  |j                  d      | _        | j                          y )NFr   )
r*   r+   r<  r.  rA  r   r   r2   r[  rI  )r1   rR   r3   s     r4   r+   zLfm2ForCausalLM.__init__=  sU     v&
 ++yy!3!3V5F5FUS 	r5   NrJ  r   r}   r   rK  labelsrL  logits_to_keepr   r(   c	           
      x    | j                   d||||||d|	}
|
j                  }t        |t              rt	        | d      n|}| j                  |dd|ddf         }d}|* | j                  d||| j                  j                  d|	}t        |||
j                  |
j                  |
j                        S )a  
        Example:

        ```python
        >>> from transformers import AutoTokenizer, Lfm2ForCausalLM

        >>> model = Lfm2ForCausalLM.from_pretrained("meta-lfm2/Lfm2-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-lfm2/Lfm2-2-7b-hf")

        >>> 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."
        ```)rJ  r   r}   r   rK  rL  N)r]  r_  rA  )lossr]  r   r6   r/  r'  )r.  rN  ru   r   slicer[  loss_functionrR   rA  r   r   r6   r/  )r1   rJ  r   r}   r   rK  r_  rL  r`  r   outputsr6   slice_indicesr]  rb  s                  r4   rC   zLfm2ForCausalLM.forwardF  s    > ,64:: ,
)%+',
 ,
  118B>SV8W~ot4]kmA}a,?@A%4%%pVFt{{OeOepiopD%#33!//))
 	
r5   )NNNNNNNr   )rH   rI   rJ   _tied_weights_keys_tp_plan_pp_planr+   r   r   r-   r+  rL   r   rW  rX  r   r   r   r   rC   rM   rN   s   @r4   rZ  rZ  7  s   *,GH23H_-z:;H  .2.204(,26*.!%-.6
##d*6
 t+6
 &&-	6

 6
 ((4/6
   4'6
 $;6
 ell*6
 +,6
 
 6
  6
r5   rZ  )rZ  r<  r-  r  )r   )Kcollections.abcr   typingr   r-   torch.nn.functionalr   r   r   cache_utilsr   r   
generationr	   integrationsr
   r   r   integrations.accelerater   masking_utilsr   r   modeling_layersr   modeling_outputsr   r   modeling_rope_utilsr   r   modeling_utilsr   r   processing_utilsr   utilsr   r   r   utils.genericr   r   utils.import_utilsr   r   utils.output_capturingr   configuration_lfm2r!   causal_conv1dr"   r#   Moduler&   rP   r   r   r   rL   r   r   rK   r   r   r   kernel_modulesallr  r   r  r-  r<  rZ  __all__r'  r5   r4   <module>r     s3  ( %      . ) f f = P 9 O K F & I I G V 5 * DD-7** Y'J")) J (J(><")) ><B8bii 8(( *+ ,2	UU\\ 	U# 	U%,, 	U& %II%<<% 
% <<	%
 LL4'% % % '(%2 )*7$BII 7$ +7$t	 #$89^, vbBII vbr*1 *Z /  & L
# L
 L
^ F
)? F
 F
R Br5   