Ë
    S^(hx  ã                   óˆ  — d dl mZmZmZ d dlmZ d dlZddlm	Z	 ddl
mZ ddlmZ ddlmZ ddlmZmZ d	d
lmZmZmZmZ d	dlmZ ddlmZ ddlmZ  ej<                  e«      Z dZ! G d„ de«      Z" G d„ dejF                  «      Z$ G d„ de«      Z% G d„ dee«      Z& G d„ de«      Z' G d„ de«      Z( G d„ de«      Z)g d¢Z*y)é    )ÚOptionalÚTupleÚUnionNé   )ÚCache)ÚFlashAttentionKwargs)ÚCausalLMOutputWithPast)ÚUnpack)Ú
LossKwargsÚloggingé   )ÚGlmAttentionÚGlmForCausalLMÚGlmForSequenceClassificationÚGlmForTokenClassification)ÚPhi3MLPé   )Ú
Glm4Config)ÚGlm4RMSNormzTHUDM/GLM-4-9B-Chat-0414c                   ó   — e Zd Zy)ÚGlm4MLPN©Ú__name__Ú
__module__Ú__qualname__© ó    úc/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/transformers/models/glm4/modular_glm4.pyr   r   *   ó   „ Ør   r   c                   óp  ‡ — e Zd Zdedefˆ fd„Z	 	 	 	 	 	 	 ddej                  deej                     deej                     dee
   dee   d	ee   d
eej                     deeej                  ej                  f      dee   deej                  eeej                  ej                  f      f   fd„Zˆ xZS )ÚGlm4DecoderLayerÚconfigÚ	layer_idxc                 ó¸  •— t         ‰| �  «        |j                  | _        t        ||¬«      | _        t        |«      | _        t        |j                  |j                  ¬«      | _	        t        |j                  |j                  ¬«      | _
        t        |j                  |j                  ¬«      | _        t        |j                  |j                  ¬«      | _        y )N)r"   r#   )Úeps)ÚsuperÚ__init__Úhidden_sizeÚGlm4AttentionÚ	self_attnr   Úmlpr   Úrms_norm_epsÚinput_layernormÚpost_attention_layernormÚpost_self_attn_layernormÚpost_mlp_layernorm)Úselfr"   r#   Ú	__class__s      €r   r'   zGlm4DecoderLayer.__init__/   s¡   ø€ Ü‰ÑÔØ!×-Ñ-ˆÔÜ&¨fÀ	ÔJˆŒä˜6“?ˆŒÜ*¨6×+=Ñ+=À6×CVÑCVÔWˆÔÜ(3°F×4FÑ4FÈF×L_ÑL_Ô(`ˆÔ%Ü(3°F×4FÑ4FÈF×L_ÑL_Ô(`ˆÔ%Ü"-¨f×.@Ñ.@Àf×FYÑFYÔ"ZˆÕr   Úhidden_statesÚattention_maskÚposition_idsÚpast_key_valueÚoutput_attentionsÚ	use_cacheÚcache_positionÚposition_embeddingsÚkwargsÚreturnc	                 ó  — |}
| j                  |«      } | j                  d||||||||dœ|	¤Ž\  }}| j                  |«      }|
|z   }|}
| j                  |«      }| j	                  |«      }| j                  |«      }|
|z   }|f}|r||fz  }|S )N)r3   r4   r5   r6   r7   r8   r9   r:   r   )r-   r*   r/   r.   r+   r0   )r1   r3   r4   r5   r6   r7   r8   r9   r:   r;   ÚresidualÚself_attn_weightsÚoutputss                r   ÚforwardzGlm4DecoderLayer.forward:   sÑ   € ð !ˆà×,Ñ,¨]Ó;ˆð ,:¨4¯>©>ð 
,
Ø'Ø)Ø%Ø)Ø/ØØ)Ø 3ñ
,
ð ñ
,
Ñ(ˆÐ(ð ×5Ñ5°mÓDˆØ  =Ñ0ˆð !ˆØ×5Ñ5°mÓDˆØŸ™ Ó/ˆØ×/Ñ/°Ó>ˆØ  =Ñ0ˆà Ð"ˆÙØÐ)Ð+Ñ+ˆGàˆr   )NNNFFNN)r   r   r   r   Úintr'   ÚtorchÚTensorr   Ú
LongTensorr   Úboolr   r
   r   ÚFloatTensorrA   Ú__classcell__©r2   s   @r   r!   r!   .   s  ø„ ð	[˜zð 	[°cõ 	[ð 26Ø37Ø*.Ø,1Ø$)Ø59ØKOñ+à—|‘|ð+ð ! §¡Ñ.ð+ð ˜u×/Ñ/Ñ0ð	+ð
 ! ™ð+ð $ D™>ð+ð ˜D‘>ð+ð ! ×!1Ñ!1Ñ2ð+ð & e¨E¯L©L¸%¿,¹,Ð,FÑ&GÑHð+ð Ð-Ñ.ð+ð 
ˆu× Ñ  (¨5°×1BÑ1BÀE×DUÑDUÐ1UÑ+VÑ"WÐWÑ	X÷+r   r!   c                   ó   — e Zd Zy)r)   Nr   r   r   r   r)   r)   h   r   r   r)   c                   ó   — e Zd Zy)ÚKwargsForCausalLMNr   r   r   r   rL   rL   l   s   … r   rL   c                   ó8   ‡ — e Zd Zdee   deeef   fˆ fd„Zˆ xZ	S )ÚGlm4ForCausalLMÚsuper_kwargsr<   c                 ó"   •— t        ‰| �  di |¤ŽS )a(  
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

            logits_to_keep (`int` or `torch.Tensor`, *optional*):
                If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
                `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
                token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
                If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
                This is useful when using packed tensor format (single dimension for batch and sequence length).

        Returns:

        Example:

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

        >>> model = Glm4ForCausalLM.from_pretrained("THUDM/GLM-4-9B-Chat-0414")
        >>> tokenizer = AutoTokenizer.from_pretrained("THUDM/GLM-4-9B-Chat-0414")

        >>> 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."
        ```r   )r&   rA   )r1   rO   r2   s     €r   rA   zGlm4ForCausalLM.forwardp   s   ø€ ôF ‰w‰Ñ. Ñ.Ð.r   )
r   r   r   r
   rL   r   r   r	   rA   rH   rI   s   @r   rN   rN   o   s0   ø„ ð#/àÐ0Ñ1ð#/ð 
ˆuÐ,Ð,Ñ	-÷#/ñ #/r   rN   c                   ó   — e Zd Zy)ÚGlm4ForSequenceClassificationNr   r   r   r   rR   rR   –   r   r   rR   c                   ó   — e Zd Zy)ÚGlm4ForTokenClassificationNr   r   r   r   rT   rT   š   r   r   rT   )ÚGlm4PreTrainedModelÚ	Glm4ModelrN   rR   rT   )+Útypingr   r   r   Útorch.nnÚnnÚtorch.utils.checkpointrC   Úcache_utilsr   Úmodeling_flash_attention_utilsr   Úmodeling_outputsr	   Úprocessing_utilsr
   Úutilsr   r   Úglm.modeling_glmr   r   r   r   Úphi3.modeling_phi3r   Úconfiguration_glm4r   Úmodeling_glm4r   Ú
get_loggerr   ÚloggerÚ_CHECKPOINT_FOR_DOCr   ÚModuler!   r)   rL   rN   rR   rT   Ú__all__r   r   r   ú<module>ri      sµ   ð÷  *Ñ )å Û å  Ý BÝ 6Ý &ß (÷ó õ )Ý *Ý &ð 
ˆ×	Ñ	˜HÓ	%€à0Ð ô	ˆgô 	ô7�r—y‘yô 7ôt	�Lô 	ô ?Ð,¨jÔ >ô$/�nô $/ôN	Ð$@ô 	ô	Ð!:ô 	ò�r   