ktj2ddlmZddlZddlmZmZddlmZmZmZm Z ddl m Z ee Z GddZGdd ZdS) ) getLoggerN) array_equalndarray) NodeProto TensorProtohelper numpy_helper) OnnxModelc VeZdZdefdZdedeeeffdZd defdZ d!ded e d edzd edzfd Z defd Z defdZ edZed"defdZdededzfdZed#defdZedefdZed$dedefdZde fdZdZdZdZdZdS)% FusionUtilsmodelc||_dSN)r )selfr s c/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/onnxruntime/transformers/fusion_utils.py__init__zFusionUtils.__init__s % input_namereturnc8|j|}|Y|jjjt jkr:||\}}t d|dd|fSt d|d|dud|fS)NzCasted graph input z to int32TzDid not cast graph input z to int32: found F) r find_graph_inputtype tensor_type elem_typerINT32cast_input_to_int32loggerdebug)rr graph_input cast_output cast_nodes rcast_graph_input_to_int32z%FusionUtils.cast_graph_input_to_int32sj11*==  "{'7'C'MQ\Qb'b'b%)%=%=j%I%I "K LLDzDDD E E E$ $ gggkaeNegghhhj  rint32c&|dz|z}|dkrttj}nO|dkrttj}n/|dkrttj}nt d||||}||fS)N_r#float32float16z"Invalid target_type: {target_type})intrrFLOATFLOAT16 ValueError add_cast_node)rr target_type output_nameto_typer!s r cast_inputzFusionUtils.cast_inputs 3&4 ' ! !++,,GG I % %++,,GG I % %+-..GGABB B&&z7KHH I%%rNr/r. graph_namech||d|z}|g}||j}||vr#||}|r|jdkr|jdg}t jd||g}|jt jd|g|j |||S)N _cast_to_Castr)inputsoutputsto)r1) r output_name_to_nodeop_typeinputr make_node attributeextendmake_attributeadd_node) rrr/r.r8r1r5 parent_noder!s rr,zFusionUtils.add_cast_node-s  $'<7'<'<>#J/ S SD|v%% >Cx4''CES9J5K5K,K,K#'S"&+a.KJ**4000J99+zRRR S Src>d}|j||vr[|||j|vrF||j||t||j|}||j|<||vr|||n|g||<|S)Nr)r:removelenappend)rJrEnew_input_namerCold_input_references rupdate_node_inputzFusionUtils.update_node_inputZs JqM0 0 0d>QRVR\]^R_>`6`6`  1 . 5 5d ; ; ;"%&9$*Q-&H"I"I & 1 0 0 0  / 6 6t < < < <37&  /""rrc|j|}|j|}t||||}|dko|| } | S)a Before: (input)-->parent-->node-->(output) After: (input)-->parent--> | +----->node-->(output) This function returns a flag whether the parent node can be removed. r)r:r rTfind_graph_output) r rJr@rCnode_input_indexparent_input_indexold_input_namerRrSparent_can_be_removeds r skip_parentzFusionUtils.skip_parentjsf$45$*+=>);;DBRTbdwxx"5!9 j5CZCZ[iCjCj?j$$rrJc|jdvsJt|jdkr%|j|jdSd}|jD]!}|jdkrtj|}"|S)N)Squeeze Unsqueezeaxes) r9rPr:r get_constant_valuer<rDrget_attribute_value)rrJr`attrs rget_squeeze_or_unsqueeze_axesz)FusionUtils.get_squeeze_or_unsqueeze_axess|77777 tz??Q  :00A?? ?N 8 8DyF""1$77 rattribute_namec|}|jD]!}|j|krtj|}"t |t r.t |t t fot||dS||kS)aVerify that a node has expected value for an attribute. Args: node (NodeProto): a node to check attribute_name (str): name of attribute expected_value (Any): expected value of the attribute default_value (Any, optional): default value if the attribute does not exist. Defaults to None. Returns: bool: whether the check is passed or not F equal_nan)r<rDrrb isinstancelistrr)rJreexpected_value default_valuevaluercs rcheck_node_attributez FusionUtils.check_node_attributesN 9 9DyN**2488 nd + + +uwo66oKX]inQuantizeLinearDequantizeLinearz+Provided node is not a Q/DQ node. Op Type: r_NFrrqT) r9rrrar:ndimshaperPryall)rJr "allow_per_tensor_quantization_onlyscalescale_has_single_element zero_points rcheck_qdq_node_for_fusionz%FusionUtils.check_qdq_node_for_fusions <E E E LLUt|UU V V V((A77 =5$):?#_uzQ7^5;WX>]^K^ - 6N 5 tz??a  4--djm<<  : ( (5  5yq)))r input_indexct|j|ksJ|j|j|}t |t r.t |t t fot||dS||kS)a7Verify that a node has expected input value Args: node (NodeProto): a node to check input_index (int): index of its input to be verified expected_value (Any): expected value of the input Returns: bool: whether the check is passed or not Frg)rPr:r rarirjrr)rrJrrkrms rcheck_node_input_valuez"FusionUtils.check_node_input_values4:,,,, --dj.EFF nd + + +uwo66oKX]inRemove Identity nodes, except those right before graph output.IdentityrzRemoved z Identity nodesN) r get_graphs_output_namesrIr9rFrHr:rQ remove_nodesrinforP)rnodes_to_removegraph_output_namesrJs rremove_identity_nodesz!FusionUtils.remove_identity_nodess!Z??AAJ$$&& 1 1D|z));q>);;;J99$+a.$*UV-XXX#**4000  J J # #O 4 4 4 KKH3#7#7HHH I I I I I J Jrc8|jdSr)r remove_cascaded_cast_nodesrs rrz&FusionUtils.remove_cascaded_cast_nodes s --/////rc8|jdSr)r remove_useless_cast_nodesrs rrz%FusionUtils.remove_useless_cast_nodes s ,,.....rcP|jd}|dSg}|jD]}|jdkr||jd}||jd}|rB|r@||kr:td|j d|| ||rTt|j }t|j }|D]}tt|j|zrtt|j|zsmt|j|jddkr2|j|jd|jdn2|j|jd|jd|j|dSdS) ziRemove reshape node that is not needed based on symbolic shape inference: input and output has same shapeT)updateNReshaperzRemove reshape node z* since its input shape is same as output: r_)r infer_runtime_shaperIr9get_edge_shaper:rFrrrDrQsetget_graphs_input_namesrboolrPrCreplace_output_of_all_nodesrHrG)r shape_inferrrJ input_shape output_shapegraph_input_namesrs rremove_useless_reshape_nodesz(FusionUtils.remove_useless_reshape_nodessj44D4AA   FJ$$&& 1 1D|y(()88AGG *99$+a.II 1<1K<4O4OKKqtyqqdoqq$**4000  - #DJ$E$E$G$G H H !$TZ%G%G%I%I!J!J ' - -DK((+==>> Y TZ3D!DEE! > > @ @A OPPTUUU >>tz!}dkZ[n]]]] J99$+a.$*UV-XXX &&t,,,, - - - -r)r#)NNN)rrr)T)__name__ __module__ __qualname__r rstrtuplerr"r0r(r,rrM staticmethodrTr[rrrdrnrrrrrrrrrrr r sU&i&&&&!C!E$)rs  &&&&&&&&============ 8  _-_-_-_-_-_-_-_-D ----------r