§
    kŠtj//  ã                   ó|   — d dl mZ d dlZd dlmZ d dlmZmZm	Z	m
Z
 d dlmZ  ee¦  «        Z G d„ de¦  «        ZdS )é    )Ú	getLoggerN)ÚFusion)Ú	NodeProtoÚTensorProtoÚhelperÚnumpy_helper)Ú	OnnxModelc                   ó”   ‡ — e Zd ZdZdededefˆ fd„Zdededeeef         fd	„Z	d
edededededededede
de
dedz  fd„Zd„ Zˆ xZS )ÚFusionAttentionVaezI
    Fuse Attention subgraph of Vae Decoder into one Attention node.
    ÚmodelÚhidden_sizeÚ	num_headsc                 óˆ   •— t          ¦   «                              |ddg¦  «         || _        || _        d| _        d| _        d S )NÚ	AttentionÚSoftmaxT)ÚsuperÚ__init__r   r   Únum_heads_warningÚhidden_size_warning)Úselfr   r   r   Ú	__class__s       €úk/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/onnxruntime/transformers/fusion_attention_vae.pyr   zFusionAttentionVae.__init__   sH   ø€ Ý‰Œ×Ò˜ ¨i¨[Ñ9Ô9Ð9Ø&ˆÔØ"ˆŒð "&ˆÔØ#'ˆÔ Ð Ð ó    Ú	reshape_qÚadd_qÚreturnc                 ó\  — | j                              |d¦  «        }|�t          |j        ¦  «        dk    r| j        | j        fS | j                              |j        d         ¦  «        }|�%t          |t          j	        ¦  «        r|j
        dk    s| j        | j        fS t          |¦  «        }|dk    r| j        | j        fS | j                              |¦  «        \  }}|�%t          |t          j	        ¦  «        r|j        dk    r| j        | j        fS |j        d         }| j        dk    r:|| j        k    r/| j        r(t                                d|| j        ¦  «         d| _        | j        dk    r:|| j        k    r/| j        r(t                                d|| j        ¦  «         d| _        ||fS )	zúDetect num_heads and hidden_size from a reshape node.

        Args:
            reshape_q (NodeProto): reshape node for Q
            add_q (NodeProto): add node for Q

        Returns:
            Tuple[int, int]: num_heads and hidden_size
        é   Né   é   r   z?Detected number of attention heads is %d. Ignore --num_heads %dFz3Detected hidden size is %d. Ignore --hidden_size %d)r   Ú
get_parentÚlenÚinputr   r   Úget_constant_valueÚ
isinstanceÚnpÚndarrayÚsizeÚintÚget_constant_inputÚndimÚshaper   ÚloggerÚwarningr   )	r   r   r   ÚconcatÚvaluer   Ú_Úbiasr   s	            r   Úget_num_heads_and_hidden_sizez0FusionAttentionVae.get_num_heads_and_hidden_size   s®  € ð ”×&Ò& y°!Ñ4Ô4ˆØˆ>�S ¤Ñ.Ô.°!Ò3Ð3Ø”> 4Ô#3Ð3Ð3à”
×-Ò-¨f¬l¸1¬oÑ>Ô>ˆØÐ!¥j°½¼
Ñ&CÔ&CÐ!ÈÌ
ÐVWÊÈØ”> 4Ô#3Ð3Ð3Ý˜‘J”Jˆ	Ø˜Š>ˆ>Ø”> 4Ô#3Ð3Ð3à”*×/Ò/°Ñ6Ô6‰ˆˆ4ØˆL¥*¨Tµ2´:Ñ">Ô">ˆLÀ4Ä9ÐPQÂ>À>Ø”> 4Ô#3Ð3Ð3à”j ”mˆàŒ>˜AÒÐ )¨t¬~Ò"=Ð"=ØÔ%ð /Ý—’ØUÐW`ÐbfÔbpñô ð ð */�Ô&àÔ˜aÒÐ K°4Ô3CÒ$CÐ$CØÔ'ð 1Ý—’ÐTÐVaÐcgÔcsÑtÔtÐtØ+0�Ô(à˜+Ð%Ð%r   Úq_matmulÚq_addÚk_matmulÚk_addÚv_matmulÚv_addÚ
input_nameÚoutput_nameNc                 ó¬	  — |j         d         |	k    s"|j         d         |	k    s|j         d         |	k    r@t                               d|j         d         |j         d         |j         d         ¦  «         dS |dk    r'||z  dk    rt                               d||¦  «         dS | j                             |j         d         ¦  «        }| j                             |j         d         ¦  «        }| j                             |j         d         ¦  «        }|r|r|sdS | j                             |j         d         ¦  «        p$| j                             |j         d         ¦  «        }| j                             |j         d         ¦  «        p$| j                             |j         d         ¦  «        }| j                             |j         d         ¦  «        p$| j                             |j         d         ¦  «        }t          j        |¦  «        }t          j        |¦  «        }t          j        |¦  «        }t          j        |j	        ¦  «        }t          j        |j	        ¦  «        }t          j        |j	        ¦  «        }|j
        dk    rt                               d¦  «         dS t          j        |¦  «        }t          j        |¦  «        }t          j        |¦  «        }|j	        |j	        k    s|j	        |j	        k    rdS |j	        d         }|j	        d         }|j	        d         }||k    r||k    sJ ‚|dk    r||k    rt          d|› d	|› d
�¦  «        ‚t          j        |j	        dd…         ¦  «        }t          j        |||fd¬¦  «        }dt          |¦  «        z  }| j                             d¦  «        } ||cxk    r|k    sn J ‚d}!t          j        |||fd¬¦  «        }"d|z  }!|                      | dz   t           j        ||g|¬¦  «         t          j        d|gt          j        ¬¦  «        }"d|z  }!|                      | dz   t           j        |!g|"¬¦  «         |	| dz   | dz   g}#t)          j        d|#|
g| ¬¦  «        }$d|$_        |$j                             t)          j        d|¦  «        g¦  «         |                      d¦  «         |$S )at  Create an Attention node.

        Args:
            q_matmul (NodeProto): MatMul node in fully connection for Q
            q_add (NodeProto): Add bias node in fully connection for Q
            k_matmul (NodeProto): MatMul node in fully connection for K
            k_add (NodeProto): Add bias node in fully connection for K
            v_matmul (NodeProto): MatMul node in fully connection for V
            v_add (NodeProto): Add bias node in fully connection for V
            num_heads (int): number of attention heads. If a model is pruned, it is the number of heads after pruning.
            hidden_size (int): hidden dimension. If a model is pruned, it is the hidden dimension after pruning.
            input_name (str): input name
            output_name (str): output name

        Returns:
            Union[NodeProto, None]: the node created or None if failed.
        r   zRFor self attention, input hidden state for q and k/v shall be same. Got %s, %s, %sNz9input hidden size %d is not a multiple of num of heads %dr   é
   zBweights are in fp16. Please run fp16 conversion after optimizationzInput hidden size (z,) is not same as weight dimension of q,k,v (z:). Please provide a correct input hidden size or pass in 0)Úaxisé   r   Ú_qkv_weight)ÚnameÚ	data_typeÚdimsÚvals)ÚdtypeÚ	_qkv_bias)ÚinputsÚoutputsrA   zcom.microsoftr   zAttention (self attention))r#   r-   Údebugr   Úget_initializerr   Úto_arrayr&   Úprodr,   rB   Ú
ValueErrorÚstackr)   Úcreate_node_nameÚadd_initializerr   ÚFLOATÚzerosÚfloat32r   Ú	make_nodeÚdomainÚ	attributeÚextendÚmake_attributeÚincrease_counter)%r   r4   r5   r6   r7   r8   r9   r   r   r:   r;   Úq_weight_tensorÚk_weight_tensorÚv_weight_tensorÚq_bias_tensorÚk_bias_tensorÚv_bias_tensorÚq_biasÚk_biasÚv_biasÚq_bias_shapeÚk_bias_shapeÚv_bias_shapeÚq_weightÚk_weightÚv_weightÚ
qw_in_sizeÚ
kw_in_sizeÚ
vw_in_sizeÚqw_out_sizeÚ
qkv_weightÚqkv_weight_dimÚattention_node_nameÚqkv_bias_dimÚqkv_biasÚattention_inputsÚattention_nodes%                                        r   Úcreate_attention_nodez(FusionAttentionVae.create_attention_nodeF   s  € ð< Œ>˜!Ô 
Ò*Ð*¨h¬n¸QÔ.?À:Ò.MÐ.MÐQYÔQ_Ð`aÔQbÐfpÒQpÐQpÝ�LŠLØdØ”˜qÔ!Ø”˜qÔ!Ø”˜qÔ!ñ	ô ð ð �4à˜Š?ˆ? ¨iÑ 7¸AÒ=Ð=Ý�LŠLÐTÐVaÐclÑmÔmÐmØ�4àœ*×4Ò4°X´^ÀAÔ5FÑGÔGˆØœ*×4Ò4°X´^ÀAÔ5FÑGÔGˆØœ*×4Ò4°X´^ÀAÔ5FÑGÔGˆØð 	 Oð 	¸ð 	Ø�4àœ
×2Ò2°5´;¸q´>ÑBÔBÐpÀdÄj×F`ÒF`ÐafÔalÐmnÔaoÑFpÔFpˆØœ
×2Ò2°5´;¸q´>ÑBÔBÐpÀdÄj×F`ÒF`ÐafÔalÐmnÔaoÑFpÔFpˆØœ
×2Ò2°5´;¸q´>ÑBÔBÐpÀdÄj×F`ÒF`ÐafÔalÐmnÔaoÑFpÔFpˆåÔ& }Ñ5Ô5ˆÝÔ& }Ñ5Ô5ˆÝÔ& }Ñ5Ô5ˆå”w˜vœ|Ñ,Ô,ˆÝ”w˜vœ|Ñ,Ô,ˆÝ”w˜vœ|Ñ,Ô,ˆð Ô$¨Ò*Ð*Ý�LŠLÐ]Ñ^Ô^Ð^Ø�4åÔ(¨Ñ9Ô9ˆÝÔ(¨Ñ9Ô9ˆÝÔ(¨Ñ9Ô9ˆð Œ>˜Xœ^Ò+Ð+¨x¬~ÀÄÒ/OÐ/OØ�4à”^ AÔ&ˆ
Ø”^ AÔ&ˆ
Ø”^ AÔ&ˆ
à˜ZÒ'Ð'¨J¸*Ò,DÐ,DÐ,DÐDà˜Š?ˆ?˜{¨jÒ8Ð8ÝðJ kð Jð JÐ_ið Jð Jð Jñô ð õ ”g˜hœn¨Q¨R¨RÔ0Ñ1Ô1ˆå”X˜x¨°8Ð<À1ÐEÑEÔEˆ
Ø�S Ñ-Ô-Ñ-ˆà"œj×9Ò9¸+ÑFÔFÐà˜|Ð;Ð;Ò;Ð;¨|Ò;Ð;Ð;Ð;Ð;Ð;àˆÝ”8˜V V¨VÐ4¸1Ð=Ñ=Ô=ˆØ˜<Ñ'ˆà×ÒØ$ }Ñ4Ý!Ô'Ø˜nÐ-Øð	 	ñ 	
ô 	
ð 	
õ ”8˜Q Ð,µB´JÐ?Ñ?Ô?ˆØ˜;‘ˆà×ÒØ$ {Ñ2Ý!Ô'Ø�Øð	 	ñ 	
ô 	
ð 	
ð Ø -Ñ/Ø +Ñ-ð
Ðõ  Ô)ØØ#Ø �MØ$ð	
ñ 
ô 
ˆð !0ˆÔØÔ ×'Ò'­Ô)>¸{ÈIÑ)VÔ)VÐ(WÑXÔXÐXà×ÒÐ:Ñ;Ô;Ð;ØÐr   c                 ó   — | j                              |d|d¬¦  «        }|€d S | j                              |d|d¬¦  «        }|€d S | j                              |d|d¬¦  «        }|€d S | j                              |d|d¬¦  «        }|€d S | j                              |d|d¬¦  «        }|€d S | j                              |d|d¬¦  «        }	|	€d S | j                              |	d|d¬¦  «        }
|
€d S | j                              |g d¢g d¢¦  «        }|€t                               d	¦  «         d S |\  }}}}}| j                              |g d
¢g d¢¦  «        }|�|\  }}}}nt                               d¦  «         d S | j                              |g d¢g d¢¦  «        }|€t                               d¦  «         d S |\  }}}}}| j                              |g d¢g d¢¦  «        }|€t                               d¦  «         d S |\  }}}}}}|}|                      ||¦  «        \  }}|dk    rt                               d¦  «         d S |                      |||||||||j        d         |j        d         ¦
  «
        }|€d S | j	         
                    |¦  «         | j        | j        |j        <   | j                             ||g¦  «         d| _        d S )NÚMatMulF)Ú	recursiveÚReshapeÚ	TransposeÚAdd)rx   ry   rx   rz   rv   )r   r   r   r   Nz&fuse_attention: failed to match v path)r   rz   ÚMulrv   )r   r   r   r   z'fuse_attention: failed to match qk path)r   r   r   r   Nz&fuse_attention: failed to match q path)ry   rx   ry   rx   rz   rv   )r   r   r   r   r   Nz&fuse_attention: failed to match k pathr   z*fuse_attention: failed to detect num_headsT)r   Úfind_first_child_by_typeÚmatch_parent_pathr-   rI   r3   rt   r#   ÚoutputÚnodes_to_addÚappendÚthis_graph_nameÚnode_name_to_graph_namerA   Únodes_to_removerW   Úprune_graph) r   Úsoftmax_nodeÚinput_name_to_nodesÚoutput_name_to_nodeÚ
matmul_qkvÚreshape_qkvÚtranspose_qkvÚreshape_outÚ
matmul_outÚadd_outÚtranspose_outÚv_nodesr1   Úadd_vÚmatmul_vÚqk_nodesÚ_softmax_qkÚ	_add_zeroÚ_mul_qkÚ	matmul_qkÚq_nodesÚ_transpose_qr   r   Úmatmul_qÚk_nodesÚadd_kÚmatmul_kÚattention_last_nodeÚq_num_headsÚq_hidden_sizeÚnew_nodes                                    r   ÚfusezFusionAttentionVae.fuseÐ   s‘  € Ø”Z×8Ò8¸ÀxÐQdÐpuÐ8ÑvÔvˆ
ØÐØˆFà”j×9Ò9¸*ÀiÐQdÐpuÐ9ÑvÔvˆØÐØˆFàœ
×;Ò;Ø˜Ð&9ÀUð <ñ 
ô 
ˆð Ð ØˆFà”j×9Ò9Ø˜9Ð&9ÀUð :ñ 
ô 
ˆð ÐØˆFà”Z×8Ò8¸ÀhÐPcÐotÐ8ÑuÔuˆ
ØÐØˆFà”*×5Ò5°jÀ%ÐI\ÐhmÐ5ÑnÔnˆØˆ?ØˆFàœ
×;Ò;¸GÀ[ÐReÐqvÐ;ÑwÔwˆØÐ ØˆFà”*×.Ò.ØÐLÐLÐLÐN`ÐN`ÐN`ñ
ô 
ˆð ˆ?Ý�LŠLÐAÑBÔBÐBØˆFØ%,Ñ"ˆˆAˆq�%˜à”:×/Ò/°
Ð<_Ð<_Ð<_ÐamÐamÐamÑnÔnˆØÐØ;CÑ8ˆ[˜) W¨i¨iå�LŠLÐBÑCÔCÐCØˆFà”*×.Ò.ØÐKÐKÐKÐM_ÐM_ÐM_ñ
ô 
ˆð ˆ?Ý�LŠLÐAÑBÔBÐBØˆFØ8?Ñ5ˆˆL˜) U¨HØ”*×.Ò.ØÐXÐXÐXÐZoÐZoÐZoñ
ô 
ˆð ˆ?Ý�LŠLÐAÑBÔBÐBØˆFØ(/Ñ%ˆˆAˆq�!�U˜Hà)Ðà%)×%GÒ%GÈ	ÐSXÑ%YÔ%YÑ"ˆ�]Ø˜!ÒÐÝ�LŠLÐEÑFÔFÐFØˆFð ×-Ò-ØØØØØØØØØŒN˜1ÔØÔ& qÔ)ñ
ô 
ˆð ÐØˆFàÔ× Ò  Ñ*Ô*Ð*Ø6:Ô6JˆÔ$ X¤]Ñ3àÔ×#Ò#Ð%8¸-Ð$HÑIÔIÐIð  ˆÔÐÐr   )Ú__name__Ú
__module__Ú__qualname__Ú__doc__r	   r)   r   r   Útupler3   Ústrrt   r¡   Ú__classcell__)r   s   @r   r   r      s>  ø€ € € € € ðð ð(˜ið (°cð (Àcð (ð (ð (ð (ð (ð (ð'&°yð '&Èð '&ÐW\Ð]`ÐbeÐ]eÔWfð '&ð '&ð '&ð '&ðRHàðHð ðHð ð	Hð
 ðHð ðHð ðHð ðHð ðHð ðHð ðHð 
�TÑ	ðHð Hð Hð HðT\ ð \ ð \ ð \ ð \ ð \ ð \ r   r   )Úloggingr   Únumpyr&   Úfusion_baser   Úonnxr   r   r   r   Ú
onnx_modelr	   r¢   r-   r   © r   r   ú<module>r¯      sº   ðð
 Ð Ð Ð Ð Ð à Ð Ð Ð Ø Ð Ð Ð Ð Ð Ø =Ð =Ð =Ð =Ð =Ð =Ð =Ð =Ð =Ð =Ð =Ð =Ø  Ð  Ð  Ð  Ð  Ð  à	ˆ�8Ñ	Ô	€ð] ð ] ð ] ð ] ð ] ˜ñ ] ô ] ð ] ð ] ð ] r   