Ë
    ÝÍ:j™( ã                   óT  — d Z ddlZddlZddlZddlZddlZddlmZ ddlm	Z	 ddl
mZ ddlZddlZddl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mZmZmZm Z 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/ ddl0m1Z1m2Z2  ejf                  d«      Z4 G d„ de«      Z5d_de6e7   dz  dejp                  fd„Z9dejp                  fd„Z:dejp                  fd„Z;d`de7de<fd„Z=d`de7de<de<fd„Z>de7de<de<de%fd „Z?d!ej*                  d"efd#„Z@d!ej*                  d"efd$„ZAd!ej*                  d"efd%„ZB	 	 	 	 dad&ed'ed(e7d)eCd*eDdz  d+eDdz  fd,„ZEd-ed.efd/„ZF	 dbd!ed)eCde6e   fd0„ZGd1„ ZHd2„ ZId3„ ZJd4efd5„ZKd4ed6e<d7e<de<fd8„ZLd4efd9„ZMd:ed;e7fd<„ZNdg fd:ed=eCd>e6eC   fd?„ZOd:efd@„ZPd:ed;e7fdA„ZQ	 	 	 dcd:edBe7dCeCdDeCdEeCf
dF„ZRd4efdG„ZSd4efdH„ZTdIefdJ„ZUd`dKe7de<fdL„ZV	 d`dKe7dMe7de<de<fdN„ZWdO„ ZXe5j²                  fdejp                  dPe5fdQ„ZZdejp                  d:ee!z  dRej¶                  dSej¶                  dTeCdUeCdVe6e6eC      deDe7ef   fdW„Z\dX„ Z]	 	 dddejp                  dYe6e7   dz  dZe<fd[„Z^d_dejp                  dYe6e7   dz  fd\„Z_dede6e7   dz  dYe6e7   dz  fd]„Z)e`d^k(  r e)«        yy)faë  
This converts GPT2 or T5 model to onnx with beam search operator.

Example 1: convert gpt2 model with beam search:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx

Example 2: convert gpt2 model with beam search containing specific cuda optimizations:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx --use_gpu                       --past_present_share_buffer --use_decoder_masked_attention

Example 3: convert gpt2 model with beam search with mixed precision and enable SkipLayerNorm strict mode:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx --use_gpu -p fp16 --use_sln_strict_mode

Example 4: convert T5 model with beam search in two steps:
    python -m models.t5.convert_to_onnx -m t5-small
    python convert_generation.py -m t5-small --model_type t5                     --decoder_onnx ./onnx_models/t5-small_decoder.onnx                       --encoder_decoder_init_onnx ./onnx_models/t5-small_encoder.onnx          --output ./onnx_models/t5_small_beam_search.onnx

Example 5: convert T5 model with beam search. All in one step:
    python convert_generation.py -m t5-small --model_type t5 --output t5_small_beam_search.onnx

Example 6: convert T5 model with beam search containing specific cuda optimizations. All in one step:
    python convert_generation.py -m t5-small --model_type t5 --output t5_small_beam_search.onnx           --use_gpu --past_present_share_buffer --use_decoder_masked_attention

Example 7: convert MT5 model with external data file like mt5-base-beamsearch.onnx.data in below example.
    python convert_generation.py -m google/mt5-base --model_type mt5 --output mt5-base-beamsearch.onnx -e

Example 8: convert gpt2 model with greedy search:
    python convert_generation.py -m gpt2 --output gpt2_greedy_search.onnx --num_beams 1 --num_return_sequences 1

Example 9: convert gpt2 model with sampling:
    python convert_generation.py -m gpt2 --output gpt2_sampling.onnx --num_beams 1 --num_return_sequences 1 --top_p 0.6
é    N)ÚEnum)ÚPath)ÚAny)Ú	PrecisionÚsetup_logger)ÚNumpyHelper)Ú
GraphProtoÚ
ModelProtoÚTensorProto)Ú	OnnxModel)Ú
GPT2ConfigÚGPT2LMHeadModelÚGPT2TokenizerÚ	MT5ConfigÚMT5ForConditionalGenerationÚT5ConfigÚT5ForConditionalGenerationÚT5Tokenizer)ÚGraphOptimizationLevelÚInferenceSessionÚSessionOptionsÚget_available_providers)Úmain)ÚPRETRAINED_GPT2_MODELS)Úexport_onnx_models)ÚPRETRAINED_MT5_MODELSÚPRETRAINED_T5_MODELSÚ c                   ó   — e Zd ZdZdZdZd„ Zy)ÚGenerationTypeÚbeam_searchÚgreedy_searchÚsamplingc                 ó   — | j                   S ©N)Úvalue)Úselfs    ú€/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/onnxruntime/transformers/convert_generation.pyÚ__str__zGenerationType.__str___   s   € Ø�z‰zÐó    N)Ú__name__Ú
__module__Ú__qualname__Ú
BEAMSEARCHÚGREEDYSEARCHÚSAMPLINGr)   © r*   r(   r    r    Z   s   „ Ø€JØ"€LØ€Hór*   r    ÚargvÚreturnc                 óp  — t        j                  «       }|j                  d«      }|j                  dddt        ddj                  t        t        z   t        z   «      z   ¬«       |j                  dd	t        d
g d¢ddj                  g d¢«      z   ¬«       |j                  dd	t        t        j                  j                  dd«      d¬«       |j                  dd	t        dd¬«       |j                  dd	t        dd¬«       |j                  dd	dd¬«       |j                  d	¬«       |j                  d«      }|j                  ddt        d¬«       |j                  d d!d	t        t        j                  j                  t        j                  j                  t        j                  j                  gd"¬«       |j                  d#d$d	d%d&gd'¬(«       |j                  d)d*d	dd+¬«       |j                  d	¬,«       |j                  d-d.d	dd/¬«       |j                  d	¬0«       |j                  d1d2d	dd3¬«       |j                  d	¬4«       |j                  d5d6d	dd7¬«       |j                  d	¬8«       |j                  d9d:d	dd;¬«       |j                  d	¬<«       |j                  d=d	dd>¬«       |j                  d	¬?«       |j                  d@«      }|j                  dAd	ddB¬«       |j                  d	¬C«       |j                  dDd	ddE¬«       |j                  d	¬F«       |j                  dGd	d¬H«       |j                  d	¬I«       |j                  dJt         d	dKdL¬M«       |j                  dNd	ddO¬«       |j                  d	¬P«       |j                  dQd	ddR¬«       |j                  d	¬S«       |j                  dTd	ddU¬«       |j                  d	¬V«       |j                  dWd	ddX¬«       |j                  d	¬Y«       |j                  dZd	dd[¬«       |j                  d	¬\«       |j                  d]d	dd^¬«       |j                  d	¬_«       |j                  d`d	dda¬«       |j                  d	¬b«       |j                  dc«      }|j                  ddt         d	dedf¬M«       |j                  dgt         d	dhdi¬M«       |j                  djt         d	dkdl¬M«       |j                  dmt         d	dedn¬M«       |j                  dot"        d	dedp¬M«       |j                  dqt"        d	dedr¬M«       |j                  dst"        d	dtdu¬M«       |j                  dvt"        d	dtdw¬M«       |j                  dxt"        d	t#        dy«       dz¬M«       |j                  d{t         d	ded|¬M«       |j                  d}t"        d	d~d¬M«       |j                  d€t         d	dKd�¬M«       |j                  d‚t         d	dƒd„¬M«       |j                  d…t         d	dƒd†¬M«       |j                  d‡t         d	dƒdˆ¬M«       |j                  d‰«      }|j                  dŠd	dd‹¬«       |j                  d	¬Œ«       |j                  d�d	ddŽ¬«       |j                  d	¬�«       |j                  d�d	dd‘¬«       |j                  d	¬’«       |j                  d“d	dd”¬«       |j                  d	¬•«       |j                  d–d	dd—¬«       |j                  d	¬˜«       |j                  d™d	t         dedš¬«       |j                  d›d	ddœ¬«       |j                  d	¬�«       |j%                  | «      }|S )žzªParse arguments

    Args:
        argv (Optional[List[str]], optional): _description_. Defaults to None.

    Returns:
        argparse.Namespace: Parsed arguments.
    zInput optionsz-mú--model_name_or_pathTzEPytorch model checkpoint path, or pretrained model name in the list: ú, )ÚrequiredÚtypeÚhelpz--model_typeFÚgpt2)r:   Út5Úmt5z*Model type (default is gpt2) in the list: )r7   r8   ÚdefaultÚchoicesr9   ú--cache_dirú.Úcache_modelsz%Directory to cache pre-trained models)r7   r8   r=   r9   z--decoder_onnxr   zLPath of onnx model for decoder. Specify it when you have exported the model.z--encoder_decoder_init_onnxzgPath of ONNX model for encoder and decoder initialization. Specify it when you have exported the model.z	--verboseÚ
store_truezPrint more information)r7   Úactionr9   )ÚverbosezOutput optionsú--outputz,Output path for onnx model with beam search.z-pú--precisionzTPrecision of model to run. fp32 for full precision, fp16 for half or mixed precisionz-bú--op_block_listÚ*ÚautozÿDisable certain onnx operators when exporting model to onnx format. When using defaultvalue for gpt2 type of model fp16 precision, it will be set to ["Add", "LayerNormalization", "SkipLayerNormalization", "FastGelu"]. Other situation, it will be set to [])r7   Únargsr=   r9   z-eú--use_external_data_formatz!save external data for model > 2G)Úuse_external_data_formatz-sz--run_shape_inferencezrun shape inference)Úrun_shape_inferencez-dpvsz--disable_pad_vocab_sizez³Do not pad logits MatMul weight to be a multiple of 8 along the dimension where dim value is the vocab size. The logits MatMul may hence be of poor performance for fp16 precision.)Údisable_pad_vocab_sizez-dsgdz,--disable_separate_gpt2_decoder_for_init_runz™Do not create separate decoder subgraphs for initial and remaining runs. This does not allow for optimizations based on sequence lengths in each subgraph)Ú*disable_separate_gpt2_decoder_for_init_runz-iz--disable_shared_initializersz™do not share initializers in encoder and decoder for T5 or in the init decoder and decoder for GPT2. It will increase memory usage of t5/mt5/gpt2 models.)Údisable_shared_initializersz--encoder_decoder_initzbAdd decoder initialization to encoder for T5 model. This is legacy format that will be deprecated.)Úencoder_decoder_initz6Beam search parameters that stored in the output modelz--output_sequences_scoreszoutput sequences scores)Úoutput_sequences_scoresz--output_token_scoreszoutput token scores)Úoutput_token_scoresz--early_stopping)r7   rC   )Úearly_stoppingz--no_repeat_ngram_sizer   zNo repeat ngram size)r8   r7   r=   r9   z--vocab_maskz\Enable vocab_mask. This mask applies only to every generated token to filter some bad words.)Ú
vocab_maskz--past_present_share_bufferzWUse shared buffer for past and present, currently work for gpt2 greedy/sampling search.)Úpast_present_share_bufferz--use_decoder_masked_attentionzëUses `DecoderMaskedSelfAttention` or `DecoderMaskedMultiHeadAttention` to optimize the decoding Attention computation. Must be used with `past_present_share_buffer`. Currently, only Attention head sizes of 32, 64 and 128 are supported.)Úuse_decoder_masked_attentionz--prefix_vocab_maskzeEnable prefix_vocab_mask. This mask can be used to filter bad words in the first generated token only)Úprefix_vocab_maskz--custom_attention_maskz]Enable custom_attention_mask. This mask can be used to replace default encoder attention mask)Úcustom_attention_maskz--presence_maskz!Presence mask for custom sampling)Úpresence_maskz--seedzRandom seed for sampling op)ÚseedzYBeam search parameters not stored in the output model, for testing parity and performancez--min_lengthé   zMin sequence lengthz--max_lengthé2   zMax sequence lengthz--num_beamsé   z	Beam sizez--num_return_sequencesz&Number of return sequence <= num_beamsz--length_penaltyz<Positive. >1 to penalize and <1 to encourage short sentence.z--repetition_penaltyz-Positive. >1 to penalize and <1 to encourage.z--temperatureç      ð?z6The value used to module the next token probabilities.z--top_pzTop P for samplingz--filter_valueÚInfzFilter value for Top P samplingz--min_tokens_to_keepzAMinimum number of tokens we keep per batch example in the output.z--presence_penaltyç        z%presence penalty for custom sampling.z--customz&If 1 customized top P logic is appliedz--vocab_sizeéÿÿÿÿzIVocab_size of the underlying model used to decide the shape of vocab maskz--eos_token_idzKcustom eos_token_id for generating model with existing onnx encoder/decoderz--pad_token_idzKcustom pad_token_id for generating model with existing onnx encoder/decoderz0Other options for testing parity and performancez--use_sln_strict_modez_Enable strict mode for SLN in CUDA provider. This ensures a better accuracy but will be slower.)Úuse_sln_strict_modeú	--use_gpuz)use GPU for inference. Required for fp16.)Úuse_gpuz--disable_parityzdo not run parity test)Údisable_parityz--disable_perf_testzdo not run perf test)Údisable_perf_testz--torch_performanceztest PyTorch performance)Útorch_performancez--total_runsz4Number of times of inference for latency measurementz--save_test_dataz-save test data for onnxruntime_perf_test tool)Úsave_test_data)ÚargparseÚArgumentParserÚadd_argument_groupÚadd_argumentÚstrÚjoinr   r   r   ÚosÚpathÚset_defaultsr   ÚFLOAT32r&   ÚFLOAT16ÚintÚfloatÚ
parse_args)r2   ÚparserÚinput_groupÚoutput_groupÚmodel_groupÚbeam_parameters_groupÚ
test_groupÚargss           r(   Úparse_argumentsr   c   sñ  € ô ×$Ñ$Ó&€Fà×+Ñ+¨OÓ<€Kà×ÑØØØÜØTØ
�)‰)Ô*Ô-AÑAÔDYÑYÓ
Zñ[ð ô ð ×ÑØØÜØÚ%Ø9¸D¿I¹IÒF[Ó<\Ñ\ð ô ð ×ÑØØÜÜ—‘—‘˜S .Ó1Ø4ð ô ð ×ÑØØÜØØ[ð ô ð ×ÑØ%ØÜØØvð ô ð ×ÑØØØØ%ð	 ô ð ×Ñ ÐÔ&à×,Ñ,Ð-=Ó>€Là×ÑØØÜØ;ð	 ô ð ×ÑØØØÜÜ×!Ñ!×'Ñ'Ü×"Ñ"×(Ñ(¬)×*;Ñ*;×*AÑ*AÐBØcð ô ð ×ÑØØØØØ�ðXð ô 	ð ×ÑØØ$ØØØ0ð ô ð ×Ñ°uÐÔ=à×ÑØØØØØ"ð ô ð ×Ñ°%ÐÔ8à×ÑØØ"ØØðbð ô ð ×Ñ°UÐÔ;à×ÑØØ6ØØðGð ô ð ×ÑÈÐÔOà×ÑØØ'ØØðEð ô ð ×Ñ¸%ÐÔ@à×ÑØ ØØØqð	 ô ð ×Ñ°5ÐÔ9à×+Ñ+Ð,dÓe€Kà×ÑØ#ØØØ&ð	 ô ð ×Ñ°UÐÔ;à×ÑØØØØ"ð	 ô ð ×Ñ°ÐÔ7à×ÑÐ/¸%ÈÐÔUØ×Ñ¨EÐÔ2à×ÑØ ÜØØØ#ð ô ð ×ÑØØØØkð	 ô ð ×Ñ¨ÐÔ.à×ÑØ%ØØØfð	 ô ð ×Ñ°uÐÔ=à×ÑØ(ØØðð	 ô ð ×Ñ¸%ÐÔ@à×ÑØØØØtð	 ô ð ×Ñ¨uÐÔ5à×ÑØ!ØØØlð	 ô ð ×Ñ°5ÐÔ9à×ÑØØØØ0ð	 ô ð ×Ñ¨5ÐÔ1à×ÑØØØØ*ð	 ô ð ×Ñ %ÐÔ(à"×5Ñ5ØcóÐð ×&Ñ& ~¼CÈ%ÐYZÐavÐ&Ôwà×&Ñ& ~¼CÈ%ÐY[ÐbwÐ&Ôxà×&Ñ& }¼3ÈÐXYÐ`kÐ&Ôlà×&Ñ&Ø ÜØØØ5ð 'ô ð ×&Ñ&ØÜØØØKð 'ô ð ×&Ñ&ØÜØØØ<ð 'ô ð ×&Ñ&ØÜØØØEð 'ô ð ×&Ñ&ØÜØØØ!ð 'ô ð ×&Ñ&ØÜØÜ�u“�Ø.ð 'ô ð ×&Ñ&ØÜØØØPð 'ô ð ×&Ñ&ØÜØØØ4ð 'ô ð ×&Ñ&ØÜØØØ5ð 'ô ð ×&Ñ&ØÜØØØXð 'ô ð ×&Ñ&ØÜØØØZð 'ô ð ×&Ñ&ØÜØØØZð 'ô ð ×*Ñ*Ð+]Ó^€Jà×ÑØØØØnð	 ô ð ×Ñ°ÐÔ6à×ÑØØØØ8ð	 ô ð ×Ñ EÐÔ*à×ÑØØØØ%ð	 ô ð ×Ñ¨5ÐÔ1à×ÑØØØØ#ð	 ô ð ×Ñ¨eÐÔ4à×ÑØØØØ'ð	 ô ð ×Ñ¨eÐÔ4à×ÑØØÜØØCð ô ð ×ÑØØØØ<ð	 ô ð ×Ñ¨5ÐÔ1à×Ñ˜TÓ"€Dà€Kr*   r~   c                 ó~  — | j                   }d|d| j                  dd| j                  ddddd	g}| j                  r|j	                  d
| j                  g«       | j
                  r|j                  d«       | j                  r|j                  d«       t        | j                  «      r-|j	                  dg«       |j	                  | j                  «       | j                  t        j                  j                  k(  r| j
                  sJ d«       ‚| j                  rt        j                  d|› �«       t!        |¬«       y)zqConvert GPT-2 model to onnx

    Args:
        args (argparse.Namespace): arguments parsed from command line
    r5   rE   z--optimize_onnxrF   z--test_runsÚ1z--test_casesÚ10z--overwriter?   rd   rK   rG   zEfp16 or mixed precision model cannot run in CPU. Please add --use_gpuzarguments for convert_to_onnx:)r2   N)Úmodel_name_or_pathÚdecoder_onnxÚ	precisionÚ	cache_dirÚextendre   ÚappendrL   ÚlenÚop_block_listr   rt   r&   rD   ÚloggerÚinfoÚconvert_gpt2_to_onnx)r~   Ú
model_nameÚ	argumentss      r(   Úgpt2_to_onnxr�   ÷  s  € ð ×(Ñ(€Jð 	ØØØ×ÑØØØ�‰ØØØØØð€Ið ‡~‚~Ø×Ñ˜-¨¯©Ð8Ô9Ø‡|‚|Ø×Ñ˜Ô%Ø×$Ò$Ø×ÑÐ5Ô6ä
ˆ4×ÑÔØ×ÑÐ+Ð,Ô-Ø×Ñ˜×+Ñ+Ô,à‡~�~œ×*Ñ*×0Ñ0Ò0Ø�|Š|ÐdÐdÓdˆ|ð
 ‡|‚|Ü�‰Ð4°Y°KÐ@ÔAä˜iÖ(r*   c                 ó&  — t        | j                  | j                  t        | j                  «      j
                  | j                  | j                  | j                  t        j                  j                  k7  | j                  ddddd| j                  | j                  | j                  t        j                  j                  k(  ¬«      }t        j                  d|d   › �«       t        j                  d|d   › �«       |d   | _        |d   | _        y)	znConvert T5 model to onnx

    Args:
        args (argparse.Namespace): arguments parsed from command line
    FT)rƒ   r†   Ú
output_dirre   rL   Úoptimize_onnxr…   rD   Úuse_decoder_start_tokenÚ	overwriteÚdisable_auto_mixed_precisionÚuse_int32_inputsÚ
model_typerQ   Úforce_fp16_iozonnx model for encoder: r   zonnx model for decoder: r\   N)Úexport_t5_onnx_modelsrƒ   r†   r   ÚoutputÚparentre   rL   r…   r   rt   r&   r˜   rQ   r‹   ÚdebugÚencoder_decoder_init_onnxr„   )r~   Úpathss     r(   Ú
t5_to_onnxr    $  sà   € ô "Ø×2Ñ2Ø—.‘.Ü˜Ÿ™Ó$×+Ñ+Ø—‘Ø!%×!>Ñ!>Ø—~‘~¬×):Ñ):×)@Ñ)@Ñ@Ø—.‘.ØØ %ØØ%*ØØ—?‘?Ø!×6Ñ6Ø—~‘~¬×):Ñ):×)@Ñ)@Ñ@ô€Eô$ ‡L�LÐ+¨E°!©H¨:Ð6Ô7Ü
‡L�LÐ+¨E°!©H¨:Ð6Ô7Ø%*¨1¡X€DÔ"Ø˜a™€DÕr*   Ú	onnx_pathrL   c                 óÆ   — ddl m} t        j                  | d¬«      }|j	                  |dd¬«      }|rt        j                  || |¬«       y	t        j                  d«       y	)
zÇShape inference on an onnx file, which will be overwritten.

    Args:
        onnx_path (str): Path of onnx model
        use_external_data_format(bool): output tensors to external data or not.
    r   )ÚSymbolicShapeInferenceT©Úload_external_dataF)Ú
auto_mergeÚguess_output_rank©Úsave_as_external_dataz4Failed to run symbolic shape inference on the model.N)	Ú&onnxruntime.tools.symbolic_shape_inferr£   ÚonnxÚ
load_modelÚinfer_shapesr   Úsaver‹   Úwarning)r¡   rL   r£   ÚmodelÚouts        r(   Úshape_inferencer²   B  sQ   € õ Nä�O‰O˜I¸$Ô?€EØ
 ×
-Ñ
-¨eÀÐX]Ð
-Ó
^€CÙ
Ü�‰�s˜IÐ=UÖVä�‰ÐMÕNr*   c                 óÄ  — t        j                  | d¬«      }|j                  j                  d   j                  }t        |«      }|j                  «       }||v sJ ‚||   }|j                  dk7  ryd}|j                  |j                  d   «      }|€9|j                  |dd«      }	|	€y|j                  |	j                  d   «      }|€yd}|j                  t        j                  j                  k7  ryt        |j                   «      dk7  ry|j                   d   }
|
d	z  dk(  ryt#        j$                  |
d	z  «      d	z  }||
z
  }|j&                  r÷|rpt)        j*                  |j                   d   |ft(        j,                  ¬
«      }t)        j.                  t1        j2                  |«      |fd¬«      }||j                   d<   not)        j*                  ||j                   d   ft(        j,                  ¬
«      }t)        j.                  t1        j2                  |«      |fd¬«      }||j                   d<   |j5                  «       |_        nyt        j6                  || |¬«       y)zâPad the logits MatMul weight in the provided decoder model, which will be overwritten.

    Args:
        onnx_path (str): Path of onnx model
        use_external_data_format(bool): output tensors to external data or not.
    Tr¤   r   ÚMatMulFr\   Ú	Transposeé   é   ©Údtype©Úaxisr¨   )r«   r¬   Úgraphr›   Únamer   Úoutput_name_to_nodeÚop_typeÚget_initializerÚinputÚmatch_parentÚ	data_typer   ÚDataTypert   r‰   ÚdimsÚmathÚceilÚraw_dataÚnpÚzerosÚfloat16Úconcatenater   Úto_arrayÚtobytesr®   )r¡   rL   Údecoder_model_protoÚlogits_output_nameÚdecoder_modelr¾   Úmatmul_nodeÚpad_along_axis_1Úlogits_weightÚtranspose_before_matmulÚactual_vocab_sizeÚpadded_vocab_sizeÚpaddingÚpadding_dataÚweight_with_paddings                  r(   Úpad_weights_of_logits_matmulrÛ   T  sQ  € ô Ÿ/™/¨)ÈÔMÐà,×2Ñ2×9Ñ9¸!Ñ<×AÑAÐäÐ1Ó2€Mà'×;Ñ;Ó=ÐØÐ!4Ñ4Ð4Ð4à%Ð&8Ñ9€Kà×Ñ˜hÒ&Øð
 ÐØ!×1Ñ1°+×2CÑ2CÀAÑ2FÓG€MØÐØ"/×"<Ñ"<¸[È+ÐWXÓ"YÐà"Ð*Øà%×5Ñ5Ð6M×6SÑ6SÐTUÑ6VÓWˆàÐ Øà Ðð ×Ñ¤+×"6Ñ"6×">Ñ">Ò>Øô ˆ=×ÑÓ !Ò#Øð &×*Ñ*¨1Ñ-Ðà˜AÑ !Ò#àäŸ	™	Ð"3°aÑ"7Ó8¸1Ñ<ÐØÐ"3Ñ3€Gð ×ÒÙÜŸ8™8 ]×%7Ñ%7¸Ñ%:¸GÐ$DÌBÏJÉJÔWˆLÜ"$§.¡.´+×2FÑ2FÀ}Ó2UÐWcÐ1dÐklÔ"mÐØ$5ˆM×Ñ˜qÒ!äŸ8™8 W¨m×.@Ñ.@ÀÑ.CÐ$DÌBÏJÉJÔWˆLÜ"$§.¡.´+×2FÑ2FÀ}Ó2UÐWcÐ1dÐklÔ"mÐØ$5ˆM×Ñ˜qÑ!à!4×!<Ñ!<Ó!>ˆÕàô ‡N�NÐ&¨	ÐIaÕbØr*   Ú
model_pathre   rc   c                 ó   — t        «       }t        j                  |_        |rddgndg}|rPdt	        «       vrt        d«      ‚t        j                  d«       |r"ddi}d|i}|D �cg c]  }||v r|||   fn|‘Œ }}t        | ||¬«      }|S c c}w )a„  Create OnnxRuntime session.

    Args:
        model_path (str): onnx model path
        use_gpu (bool): use GPU or not
        use_sln_strict_mode (bool): use strict mode for skip layer normalization or not

    Raises:
        RuntimeError: CUDAExecutionProvider is not available when --use_gpu is specified.

    Returns:
        onnxruntime.InferenceSession: The created session.
    ÚCUDAExecutionProviderÚCPUExecutionProviderz5CUDAExecutionProvider is not available for --use_gpu!zuse CUDAExecutionProviderÚ"enable_skip_layer_norm_strict_modeT)Ú	providers)	r   r   ÚORT_DISABLE_ALLÚgraph_optimization_levelr   ÚRuntimeErrorr‹   rŒ   r   )	rÜ   re   rc   Úsess_optionsÚexecution_providersÚcuda_provider_optionsÚprovider_optionsr½   Úort_sessions	            r(   Úcreate_ort_sessionrê   ¢  sÁ   € ô "Ó#€LÜ,B×,RÑ,R€LÔ)ÙOVÐ2Ð4JÑKÐ]sÐ\tÐÙØ"Ô*AÓ*CÑCÜÐVÓWÐWä�K‰KÐ3Ô4ÙØ%IÈ4Ð$PÐ!Ø 7Ð9NÐOÐàatö#ØY]°$Ð:JÑ2J�Ð'¨Ñ-Ñ.ÐPTÑTð#Ðð #ô # :¨|ÐGZÔ[€KØÐùò#s   Á%Br¼   r…   c           
      óÔ  — |t         j                  j                  k(  }t        | j                  «      }|dz
  }|dk\  sJ ‚g d¢t        |«      D �cg c]  }d|› �‘Œ	 c}z   }t        | j                  «      t        |«      k7  r-t        dt        |«      › dt        | j                  «      › �«      ‚t        |«      D ]É  \  }}| j                  |   j                  |k7  r+t        d|› d|› d| j                  |   j                  › �«      ‚t        j                  }|dk\  r"|rt        j                  nt        j                  }| j                  |   j                  j                  j                  }	|	|k7  sŒ·t        d|› d	|› d|	› �«      ‚ t        j!                  d
«       dgt        |«      D �cg c]  }d|› �‘Œ	 c}z   }
t        | j"                  «      t        |
«      k7  r-t        dt        |
«      › dt        | j"                  «      › �«      ‚t        |
«      D ]´  \  }}| j"                  |   j                  |k7  r+t        d|› d|› d| j"                  |   j                  › �«      ‚|rt        j                  nt        j                  }| j"                  |   j                  j                  j                  }||k7  sŒ¢t        d|› d	|› d|› �«      ‚ t        j!                  d«       yc c}w c c}w )aæ  Verify GPT-2 subgraph

    Args:
        graph (onnx.GraphProto): onnx graph of GPT-2
        precision (Precision): Precision (FLOAT16 or FLOAT32) of the model.

    Raises:
        ValueError: Number of inputs not expected.
        ValueError: Input name is not expected.
        ValueError: Input data type is not expected.
        ValueError: Number of outputs not expected.
        ValueError: Output name is not expected.
        ValueError: Output data type is not expected.
    é   r\   )Ú	input_idsÚposition_idsÚattention_maskÚpast_ú Number of inputs expected to be ú. Got úInput ú is expected to be ú$ is expected to have onnx data type z:Verifying GPT-2 graph inputs: name and data type are good.ÚlogitsÚpresent_ú!Number of outputs expected to be úOutput z;Verifying GPT-2 graph outputs: name and data type are good.N)r   rt   r&   r‰   rÁ   ÚrangeÚ
ValueErrorÚ	enumerater½   r   ÚINT32ÚFLOATr8   Útensor_typeÚ	elem_typer‹   rŒ   r›   )r¼   r…   Ú
is_float16Úinput_countÚlayer_countÚiÚexpected_inputsÚexpected_inputÚexpected_typeÚ
input_typeÚexpected_outputsÚexpected_outputÚoutput_types                r(   Úverify_gpt2_subgraphr  Ã  sÛ  € ð œi×/Ñ/×5Ñ5Ñ5€Jä�e—k‘kÓ"€KØ ‘/€KØ˜!ÒÐÐâEÔ^cÐdoÓ^pÖHqÐYZÈ5ÐQRÐPSÊÒHqÑq€OÜ
ˆ5�;‰;Óœ3˜Ó/Ò/ÜÐ;¼CÀÓ<PÐ;QÐQWÔX[Ð\a×\gÑ\gÓXhÐWiÐjÓkÐkä& Ó7ò 
pÑˆˆ>Ø�;‰;�q‰>×Ñ .Ò0Ü˜v a SÐ(;¸NÐ;KÈ6ÐRW×R]ÑR]Ð^_ÑR`×ReÑReÐQfÐgÓhÐhä#×)Ñ)ˆØ�Š6Ù3=œK×/Ò/Ä;×CTÑCTˆMà—[‘[ ‘^×(Ñ(×4Ñ4×>Ñ>ˆ
Ø˜Ó&Ü˜v a SÐ(LÈ]ÈOÐ[aÐblÐamÐnÓoÐoð
pô ‡K�KÐLÔMà �z¼UÀ;Ó=OÖ$P¸ x°¨s¢^Ò$PÑPÐÜ
ˆ5�<‰<ÓœCÐ 0Ó1Ò1ÜÐ<¼SÐAQÓ=RÐ<SÐSYÔZ]Ð^c×^jÑ^jÓZkÐYlÐmÓnÐnä'Ð(8Ó9ò qÑˆˆ?Ø�<‰<˜‰?×Ñ ?Ò2Ü˜w q cÐ)<¸_Ð<MÈVÐTY×T`ÑT`ÐabÑTc×ThÑThÐSiÐjÓkÐká/9œ×+Ò+¼{×?PÑ?PˆØ—l‘l 1‘o×*Ñ*×6Ñ6×@Ñ@ˆØ˜-Ó'Ü˜v a SÐ(LÈ]ÈOÐ[aÐbmÐanÐoÓpÐpðqô ‡K�KÐMÔNð ùòA Irùò" %Qs   ÁK Æ)K%c           
      ó  — |t         j                  j                  k(  }|rt        j                  nt        j                  }t        | j                  «      }|dz
  dz  }|dk\  sJ ‚ddg}t        |«      D ]*  }|j                  d|› �«       |j                  d|› �«       Œ, t        |«      D ]*  }|j                  d|› �«       |j                  d	|› �«       Œ, t        | j                  «      t        |«      k7  r-t        d
t        |«      › dt        | j                  «      › �«      ‚t        |«      D ]©  \  }}| j                  |   j                  |k7  r+t        d|› d|› d| j                  |   j                  › �«      ‚|dk  rt        j                  n|}	| j                  |   j                  j                  j                  }
|
|	k7  sŒ—t        d|› d|	› d|
› �«      ‚ dg}t        |«      D ]*  }|j                  d|› �«       |j                  d|› �«       Œ, t        | j                   «      t        |«      k7  r-t        dt        |«      › dt        | j                   «      › �«      ‚t        |«      D ]’  \  }}| j                   |   j                  |k7  r+t        d|› d|› d| j                   |   j                  › �«      ‚| j                   |   j                  j                  j                  }||k7  sŒ€t        d|› d|› d|› �«      ‚ y)áð  Verify T5 decoder subgraph

    Args:
        graph (onnx.GraphProto): onnx graph of T5 decoder
        precision (Precision): Precision (FLOAT16 or FLOAT32) of the model.

    Raises:
        ValueError: Number of inputs not expected.
        ValueError: Input name is not expected.
        ValueError: Input data type is not expected.
        ValueError: Number of outputs not expected.
        ValueError: Output name is not expected.
        ValueError: Output data type is not expected.
    r¶   r^   r\   rí   Úencoder_attention_maskÚpast_key_self_Úpast_value_self_Úpast_key_cross_Úpast_value_cross_rñ   rò   ró   rô   rõ   rö   Úpresent_key_self_Úpresent_value_self_rø   rù   N)r   rt   r&   r   rþ   r‰   rÁ   rú   rˆ   rû   rü   r½   rý   r8   rÿ   r   r›   )r¼   r…   r  Ú
float_typer  r  r  r  r  r  r  r	  r
  r  s                 r(   Úverify_t5_decoder_subgraphr  û  s.  € ð œi×/Ñ/×5Ñ5Ñ5€JÙ(2”×$Ò$¼×8IÑ8I€Jä�e—k‘kÓ"€KØ ‘? qÑ(€KØ˜!ÒÐÐð #Ð$<Ð=€OÜ�;Óò 7ˆØ×Ñ °¨sÐ3Ô4Ø×ÑÐ!1°!°Ð5Õ6ð7ô �;Óò 8ˆØ×Ñ °°Ð4Ô5Ø×ÑÐ!2°1°#Ð6Õ7ð8ô ˆ5�;‰;Óœ3˜Ó/Ò/ÜÐ;¼CÀÓ<PÐ;QÐQWÔX[Ð\a×\gÑ\gÓXhÐWiÐjÓkÐkä& Ó7ò pÑˆˆ>Ø�;‰;�q‰>×Ñ .Ò0Ü˜v a SÐ(;¸NÐ;KÈ6ÐRW×R]ÑR]Ð^_ÑR`×ReÑReÐQfÐgÓhÐhà-.°ªUœ×)Ò)¸
ˆØ—[‘[ ‘^×(Ñ(×4Ñ4×>Ñ>ˆ
Ø˜Ó&Ü˜v a SÐ(LÈ]ÈOÐ[aÐblÐamÐnÓoÐoðpð !�zÐÜ�;Óò ;ˆØ×ÑÐ"3°A°3Ð 7Ô8Ø×ÑÐ"5°a°SÐ 9Õ:ð;ô ˆ5�<‰<ÓœCÐ 0Ó1Ò1ÜÐ<¼SÐAQÓ=RÐ<SÐSYÔZ]Ð^c×^jÑ^jÓZkÐYlÐmÓnÐnä'Ð(8Ó9ò oÑˆˆ?Ø�<‰<˜‰?×Ñ ?Ò2Ü˜w q cÐ)<¸_Ð<MÈVÐTY×T`ÑT`ÐabÑTc×ThÑThÐSiÐjÓkÐkØ—l‘l 1‘o×*Ñ*×6Ñ6×@Ñ@ˆØ˜*Ó$Ü˜w q cÐ)MÈjÈ\ÐY_Ð`kÐ_lÐmÓnÐnñor*   c           
      óX  — |t         j                  j                  k(  }d| j                  d   j                  v }g d¢}|r|dd }t        | j                  «      t        |«      k7  r-t        dt        |«      › dt        | j                  «      › �«      ‚t        |«      D ]¢  \  }}| j                  |   j                  |k7  r+t        d|› d	|› d| j                  |   j                  › �«      ‚t        j                  }| j                  |   j                  j                  j                  }||k7  sŒ�t        d|› d
|› d|› �«      ‚ |rwt        | j                  «      dz  dk(  sJ ‚t        | j                  «      dz  }	|	dk\  sJ ‚g }
t        |	«      D ]*  }|
j                  d|› �«       |
j                  d|› �«       Œ, nËt         j#                  d«       t        | j                  «      dz
  dz  dk(  sJ ‚t        | j                  «      dz
  dz  }	|	dk\  sJ ‚ddg}
t        |	«      D ]*  }|
j                  d|› �«       |
j                  d|› �«       Œ, t        |	«      D ]*  }|
j                  d|› �«       |
j                  d|› �«       Œ, t        | j                  «      t        |
«      k7  r-t        dt        |
«      › dt        | j                  «      › �«      ‚t        |
«      D ]´  \  }}| j                  |   j                  |k7  r+t        d|› d	|› d| j                  |   j                  › �«      ‚|rt        j                  nt        j$                  }| j                  |   j                  j                  j                  }||k7  sŒ¢t        d|› d
|› d|› �«      ‚ t         j'                  d«       y)r  Úcrossr   )Úencoder_input_idsr  Údecoder_input_idsNr¶   rñ   rò   ró   rô   rõ   r\   Úpresent_key_cross_Úpresent_value_cross_zZThis format is deprecated. Please export T5 encoder in new format with only cross outputs.r^   rö   Úencoder_hidden_statesr  r  rø   rù   zMT5 encoder graph verified: name and data type of inputs and outputs are good.)r   rt   r&   r›   r½   r‰   rÁ   rû   rü   r   rý   r8   rÿ   r   rú   rˆ   r‹   r¯   rþ   rŒ   )r¼   r…   r  Ú
new_formatr  r  r  r  r  r  r	  r
  r  s                r(   Ú'verify_t5_encoder_decoder_init_subgraphr   G  sÖ  € ð œi×/Ñ/×5Ñ5Ñ5€JØ˜EŸL™L¨™O×0Ñ0Ð0€Jò€Oñ
 Ø)¨"¨1Ð-ˆÜ
ˆ5�;‰;Óœ3˜Ó/Ò/ÜÐ;¼CÀÓ<PÐ;QÐQWÔX[Ð\a×\gÑ\gÓXhÐWiÐjÓkÐkä& Ó7ò pÑˆˆ>Ø�;‰;�q‰>×Ñ .Ò0Ü˜v a SÐ(;¸NÐ;KÈ6ÐRW×R]ÑR]Ð^_ÑR`×ReÑReÐQfÐgÓhÐhä#×)Ñ)ˆØ—[‘[ ‘^×(Ñ(×4Ñ4×>Ñ>ˆ
Ø˜Ó&Ü˜v a SÐ(LÈ]ÈOÐ[aÐblÐamÐnÓoÐoðpñ Ü�5—<‘<Ó  1Ñ$¨Ò)Ð)Ð)Ü˜%Ÿ,™,Ó'¨1Ñ,ˆØ˜aÒÐÐð ÐÜ�{Ó#ò 	@ˆAØ×#Ñ#Ð&8¸¸Ð$<Ô=Ø×#Ñ#Ð&:¸1¸#Ð$>Õ?ñ	@ô 	�‰ÐsÔtÜ�E—L‘LÓ! AÑ%¨Ñ*¨aÒ/Ð/Ð/Ü˜5Ÿ<™<Ó(¨1Ñ,°Ñ2ˆØ˜aÒÐÐð %Ð&=Ð>ÐÜ�{Ó#ò 	?ˆAØ×#Ñ#Ð&7¸°sÐ$;Ô<Ø×#Ñ#Ð&9¸!¸Ð$=Õ>ð	?ô �{Ó#ò 	@ˆAØ×#Ñ#Ð&8¸¸Ð$<Ô=Ø×#Ñ#Ð&:¸1¸#Ð$>Õ?ð	@ô ˆ5�<‰<ÓœCÐ 0Ó1Ò1ÜÐ<¼SÐAQÓ=RÐ<SÐSYÔZ]Ð^c×^jÑ^jÓZkÐYlÐmÓnÐnä'Ð(8Ó9ò rÑˆˆ?Ø�<‰<˜‰?×Ñ ?Ò2Ü˜w q cÐ)<¸_Ð<MÈVÐTY×T`ÑT`ÐabÑTc×ThÑThÐSiÐjÓkÐká/9œ×+Ò+¼{×?PÑ?PˆØ—l‘l 1‘o×*Ñ*×6Ñ6×@Ñ@ˆØ˜-Ó'Ü˜w q cÐ)MÈmÈ_Ð\bÐcnÐboÐpÓqÐqðrô ‡K�KÐ_Õ`r*   Úgraph1Úgraph2Úshared_prefixÚmin_elementsÚsignature_cache1Úsignature_cache2c                 ó	  — i }i }g }g }	g }
| j                   D ]ò  }|j                  rt        |j                  «      |k\  sŒ(|j                   D ]¼  }|j                  rt        |j                  «      |k\  sŒ(t        j                  ||||«      sŒA||j
                  z   ||j
                  <   |j                  |«       |j
                  |vr@||j
                  z   }|||j
                  <   |	j                  |«       |
j                  |«        Œò Œô t        j                  d|
› �«       | j                  D ]Q  }t        t        |j                  «      «      D ].  }|j                  |   |
v sŒt        d|j                  |   › �«      ‚ ŒS |j                  D ]Q  }t        t        |j                  «      «      D ].  }|j                  |   |
v sŒt        d|j                  |   › �«      ‚ ŒS |	D ]  }|j                   j                  |«       Œ |j                  D ]%  }|j
                  |v sŒ||j
                     |_        Œ' |j                  D ]�  }t        t        |j                  «      «      D ]m  }|j                  |   |v sŒ||j                  |      }t        j                  d|j
                  › d|› d|j                  |   › d|› �«       ||j                  |<   Œo Œ’ |D ]  }| j                   j                  |«       Œ | j                  D ]%  }|j
                  |v sŒ||j
                     |_        Œ' | j                  D ]�  }t        t        |j                  «      «      D ]m  }|j                  |   |v sŒ||j                  |      }t        j                  d|j
                  › d|› d|j                  |   › d|› �«       ||j                  |<   Œo Œ’ |	D ]  }||j
                     |_        Œ |	D ]–  }t         j"                  j%                  |«      j&                  }t         j(                  j+                  |j
                  |j,                  |«      }| j                  j                  |«       |j                  j                  |«       Œ˜ |	S )	a…  Remove initializers with same value from two graphs.

    Args:
        graph1 (GraphProto): the first graph to process
        graph2 (GraphProto): the second graph to process
        shared_prefix (str): add prefix to the shared initializers among two graphs
        min_elements (int, optional): minimal number of elements for initializers to be considered. Defaults to 1024.
        signature_cache1 (dict): Optional dictionary to store data signatures of tensors in graph1 in order to speed up comparison
        signature_cache2 (dict): Optional dictionary to store data signatures of tensors in graph2 in order to speed up comparison
    zshared initializers:zname is found in graph 1: zname is found in graph 2: zgraph 2 rename node z input z from z to zgraph 1 rename node )ÚinitializerrÅ   Úsumr   Úhas_same_valuer½   rˆ   r‹   r�   Únoderú   r‰   rÁ   rä   ÚremoveÚ
value_infor«   Únumpy_helperrÍ   ÚshapeÚhelperÚmake_tensor_value_inforÃ   )r!  r"  r#  r$  r%  r&  Úmapping_initializers_1Úmapping_initializers_2Úshared_initializers_1Úshared_initializers_2Úshared_initializers_namesÚinitializer1Úinitializer2Úshared_namer+  Újr(  r-  Únew_namer/  s                       r(   Úremove_shared_initializersr<  £  s}  € ð&  ÐØÐØÐØÐØ "Ðà×*Ñ*ò ˆØ×!Ò!¤c¨,×*;Ñ*;Ó&<ÀÒ&LØà"×.Ñ.ò 	ˆLØ ×%Ò%¬#¨l×.?Ñ.?Ó*@ÀLÒ*PØä×'Ñ'¨°lÐDTÐVfÕgØ<IÈL×L]ÑL]Ñ<]Ð& |×'8Ñ'8Ñ9Ø%×,Ñ,¨\Ô:à×$Ñ$Ð,BÑBØ"/°,×2CÑ2CÑ"C�KØ@KÐ*¨<×+<Ñ+<Ñ=Ø)×0Ñ0°Ô>Ø-×4Ñ4°[ÔAÙñ	ð	ô& ‡L�LÐ'Ð(AÐ'BÐCÔDð —‘ò QˆÜ”s˜4Ÿ:™:“Ó'ò 	QˆAØ�z‰z˜!‰}Ð 9Ò9Ü"Ð%?ÀÇ
Á
È1Á¸Ð#OÓPÐPñ	QðQð —‘ò QˆÜ”s˜4Ÿ:™:“Ó'ò 	QˆAØ�z‰z˜!‰}Ð 9Ò9Ü"Ð%?ÀÇ
Á
È1Á¸Ð#OÓPÐPñ	QðQð -ò /ˆØ×Ñ×!Ñ! +Õ.ð/ð ×'Ñ'ò Fˆ
Ø�?‰?Ð4Ò4Ø4°Z·_±_ÑEˆJ�OðFð
 —‘ò )ˆÜ”s˜4Ÿ:™:“Ó'ò 	)ˆAØ�z‰z˜!‰}Ð 6Ò6Ø1°$·*±*¸Q±-Ñ@�Ü—‘Ð3°D·I±I°;¸gÀaÀSÈÈtÏzÉzÐZ[É}ÈoÐ]aÐbjÐakÐlÔmØ (�—
‘
˜1’ñ		)ð)ð -ò /ˆØ×Ñ×!Ñ! +Õ.ð/ð ×'Ñ'ò Fˆ
Ø�?‰?Ð4Ò4Ø4°Z·_±_ÑEˆJ�OðFð
 —‘ò )ˆÜ”s˜4Ÿ:™:“Ó'ò 	)ˆAØ�z‰z˜!‰}Ð 6Ò6Ø1°$·*±*¸Q±-Ñ@�Ü—‘Ð3°D·I±I°;¸gÀaÀSÈÈtÏzÉzÐZ[É}ÈoÐ]aÐbjÐakÐlÔmØ (�—
‘
˜1’ñ		)ð)ð -ò DˆØ1°+×2BÑ2BÑCˆÕðDð -ò -ˆÜ×!Ñ!×*Ñ*¨;Ó7×=Ñ=ˆÜ—[‘[×7Ñ7¸×8HÑ8HÈ+×J_ÑJ_ÐafÓgˆ
à×Ñ× Ñ  Ô,Ø×Ñ× Ñ  Õ,ð-ð !Ð r*   Úencoder_modelrÑ   c                 ó2  — t        | «      }t        |«      }|j                  d«       |j                  d«       i i }}|j                  |«       |j                  |«       t        |j                  j
                  |j                  j
                  d||¬«      }|S )NÚe_Úd_Ús_)r#  r%  r&  )r   Úadd_prefix_to_namesÚremove_duplicated_initializerr<  r°   r¼   )r=  rÑ   ÚencoderÚdecoderr%  r&  Úinitializerss          r(   Úget_shared_initializersrG    s�   € Ü˜Ó&€GÜ˜Ó&€GØ×Ñ Ô%Ø×Ñ Ô%Ø)+¨RÐ&ÐØ×)Ñ)Ð*:Ô;Ø×)Ñ)Ð*:Ô;Ü-Ø�‰×ÑØ�‰×ÑØØ)Ø)ô€Lð Ðr*   c                 óÜ  — g }| j                   D ]8  }|j                  rt        |j                  «      |k\  sŒ(|j                  |«       Œ: |D ]  }| j                   j	                  |«       Œ |D ]{  }t
        j                  j                  |«      j                  }t
        j                  j                  |j                  |j                  |«      }| j                  j                  |«       Œ} |S )a^  Remove initializers of a graph, when they have number of elements larger than a threshold.

    Args:
        graph (GraphProto): the graph.
        min_elements (int, optional): minimal number of elements for initializers to be considered. Defaults to 1024.

    Returns:
        List[TensorProto]: initializers that are removed from the graph.
    )r(  rÅ   r)  rˆ   r,  r«   r.  rÍ   r/  r0  r1  r½   rÃ   r-  )r¼   r$  Úmoved_initializersÚtensorr(  r/  r-  s          r(   Úmove_initializersrK    sÚ   € ð ÐØ×#Ñ#ò *ˆØ—’¤ F§K¡KÓ 0°LÒ @ØØ×!Ñ! &Õ)ð*ð
 *ò .ˆØ×Ñ× Ñ  Õ-ð.ð *ò ,ˆÜ×!Ñ!×*Ñ*¨;Ó7×=Ñ=ˆÜ—[‘[×7Ñ7¸×8HÑ8HÈ+×J_ÑJ_ÐafÓgˆ
Ø×Ñ×Ñ 
Õ+ð,ð
 Ðr*   c                 óî  — | j                   dk(  rt        d| j                  › d�«      ‚| j                   dk(  r| j                  }�n#| j                   dk(  r| j                  }�n| j                   dk(  r| j
                  }nê| j                   dk(  r| j                  }nÎ| j                   dk(  r| j                  }n²| j                   d	k(  r| j                  }n–| j                   d
k(  r| j                  }nz| j                   dk(  r| j                  }n^| j                   dk(  r| j                  }nB| j                   dk(  r| j                  }n&t        d| j                  › d| j                   › d�«      ‚| j                  |fS )zÄ
    Convert attribute to kwarg format for use with onnx.helper.make_node.
        :parameter attribute: attribute in AttributeProto format.
        :return: attribute in {key: value} format.
    r   z
attribute z does not have type specified.r\   r¶   rì   r^   é   é   é   r·   é	   é
   z has unsupported type r@   )r8   rû   r½   Úfr  ÚsÚtÚgÚfloatsÚintsÚstringsÚtensorsÚgraphs)Ú	attributer&   s     r(   Ú_attribute_to_pairr\  >  s;  € ð ‡~�~˜ÒÜ˜: i§n¡nÐ%5Ð5SÐTÓUÐUð ‡~�~˜ÒØ—‘ŠØ	�‰˜1Ò	Ø—‘ŠØ	�‰˜1Ò	Ø—‘‰Ø	�‰˜1Ò	Ø—‘‰Ø	�‰˜1Ò	Ø—‘‰Ø	�‰˜1Ò	Ø× Ñ ‰Ø	�‰˜1Ò	Ø—‘‰Ø	�‰˜1Ò	Ø×!Ñ!‰Ø	�‰˜1Ò	Ø×!Ñ!‰Ø	�‰˜2Ò	Ø× Ñ ‰ä˜: i§n¡nÐ%5Ð5KÈIÏNÉNÐK[Ð[\Ð]Ó^Ð^à�N‰N˜EÐ"Ð"r*   c                 óÀ   — i }| j                   D ]#  }t        |«      \  }}|j                  ||i«       Œ% | j                  r|j                  d| j                  i«       |S )NÚdomain)r[  r\  Úupdater^  )r+  ÚkwargsÚattrÚkeyr&   s        r(   Ú	kwargs_ofrc  c  s[   € Ø€FØ—‘ò $ˆÜ)¨$Ó/‰ˆˆeØ�‰�s˜E�lÕ#ð$ð ‡{‚{Ø�‰�x §¡Ð-Ô.Ø€Mr*   c                 óÒ   — t        | j                  j                  j                  j                  D �cg c]&  }|j
                  r|j
                  n|j                  ‘Œ( c}«      S c c}w r%   )Útupler8   rÿ   r/  ÚdimÚ	dim_paramÚ	dim_value)ÚviÚds     r(   Úshape_ofrk  m  sD   € ÜÈÏÉ×I\ÑI\×IbÑIb×IfÑIfÖgÀA !§+¢+�!—+’+°A·K±KÑ?ÒgÓhÐhùÒgs   ²+A$Úsubgc                 ó  — d}d}g }t        | j                  «      D ]‚  \  }}||k\  rft        |«      }t        j                  j                  |j                  |j                  j                  j                  |d   |d   |d   d|d   g¬«      }|j                  |g«       Œ„ |j                  t        j                  j                  dt        j                  j                  dg¬	«      g«       | j                  d
«       | j                  j                  |«       g }t        | j                  «      D ]‚  \  }}||k\  rft        |«      }t        j                  j                  |j                  |j                  j                  j                  |d   |d   |d   d|d   g¬«      }|j                  |g«       Œ„ | j                  d«       | j                  j                  |«       g }| j                  D ]ç  }	|	}
|	j                   dk(  rÂt#        |	«      }|j%                  ddi«       g }|j                  |	j                  «       t'        |«      dk  r!|j                  dg«       t'        |«      dk  rŒ!t'        |«      dk  r|j                  dg«       t        j                  j(                  d||	j                  fd|	j                  i|¤Ž}
|j                  |
g«       Œé | j                  d«       | j                  j                  |«       | S )Nrì   r\   r   r¶   Úmax_seq_lenr^   ©r   r/  Úpast_sequence_length©r/  rÁ   r›   Ú	AttentionrV   rN  r   rO  r½   r+  )rü   rÁ   rk  r«   r0  r1  r½   r8   rÿ   r   r‡   r   rý   Ú
ClearFieldr›   r+  r¿   rc  r_  r‰   Ú	make_node)rl  Úinput_past_0Úoutput_past_0Ú
new_inputsr  ri  r/  Únew_outputsÚ	new_nodesr+  Únew_noder`  Úniss                r(   Ú1update_decoder_subgraph_past_present_share_bufferr|  q  s¬  € Ø€LØ€MØ€JÜ˜4Ÿ:™:Ó&ò  ‰ˆˆ2Ø�ÒÜ˜R“LˆEÜ—‘×3Ñ3Ø—‘ØŸ'™'×-Ñ-×7Ñ7Ø˜Q‘x  q¡¨5°©8°]ÀEÈ!ÁHÐMð 4ó ˆBð
 	×Ñ˜2˜$Õð ð ×Ñ”t—{‘{×9Ñ9Ð:PÔRV×RbÑRb×RhÑRhÐqrÐpsÐ9ÓtÐuÔvØ‡O�O�GÔØ‡J�J×Ñ�jÔ!à€KÜ˜4Ÿ;™;Ó'ò !‰ˆˆ2Ø�ÒÜ˜R“LˆEÜ—‘×3Ñ3Ø—‘ØŸ'™'×-Ñ-×7Ñ7Ø˜Q‘x  q¡¨5°©8°]ÀEÈ!ÁHÐMð 4ó ˆBð
 	×Ñ˜B˜4Õ ð!ð 	‡O�O�HÔØ‡K�K×Ñ�{Ô#à€IØ—	‘	ò %ˆØˆØ�<‰<˜;Ò&Ü˜t“_ˆFØ�M‰MÐ6¸Ð:Ô;ØˆCØ�J‰J�t—z‘zÔ"Ü�c“(˜Q’,Ø—
‘
˜B˜4Ô ô �c“(˜Q“,ä�3‹x˜!Š|Ø—
‘
Ð2Ð3Ô4Ü—{‘{×,Ñ,¨[¸#¸t¿{¹{ÑeÐQU×QZÑQZÐeÐ^dÑeˆHØ×Ñ˜(˜Õ$ð%ð 	‡O�O�FÔØ‡I�I×Ñ�YÔØ€Kr*   Úis_beam_searchÚswitch_attentionc                 óì  — |rôg }t        | j                  «      D ]  \  }}|j                  |g«       Œ |j                  t        j                  j                  dt        j                  j                  dg¬«      g«       |j                  t        j                  j                  dt        j                  j                  g d¢¬«      g«       | j                  d«       | j                  j                  |«       |�r{g d¢}g }| j                  D �]9  }|j                  dk(  �rt        |«      }	|	j                  «       D ]0  }
|
d	k(  r  y
|
|vsŒ|
dk7  rt        j                  d|
› d�«       |	|
= Œ2 g }|j                  |j                  «       |rot        |«      dk  r!|j                  dg«       t        |«      dk  rŒ!t        |«      dk  r|j                  dg«       t        |«      dk  r|j                  dg«       t        j                  j                   d||j"                  fd|j$                  i|	¤Ž}|j                  |g«       �Œ< | j                  d«       | j                  j                  |«       y)aS  Update the Attention nodes to DecoderMaskedSelfAttention.

    Args:
        subg (GraphProto): GraphProto of the decoder subgraph
        is_beam_search (bool): Boolean specifying if the sampling algo is BeamSearch
        switch_attention (bool): Boolean specifying if `Attention` is to be switched with `DecoderMaskedSelfAttention`
    Ú
beam_widthr\   rq  Úcache_indirection©Ú
batch_sizer€  rn  rÁ   ©rV   Ú	num_headsÚscaleÚmask_filter_valuer^  rr  Úqkv_hidden_sizesFÚunidirectionalzRemoving attribute: zB from Attention node while switching to DecoderMaskedSelfAttentionrO  r   r·   rP  ÚDecoderMaskedSelfAttentionr½   r+  T)rü   rÁ   r‡   r«   r0  r1  r   rý   rs  r+  r¿   rc  Úcopyr‹   r¯   r‰   rt  r›   r½   )rl  r}  r~  rw  Ú_iri  Ú'decoder_masked_attention_supported_attrry  r+  r`  Úkr{  s               r(   Ú4update_decoder_subgraph_use_decoder_masked_attentionr�  ¢  sI  € ñ Øˆ
Ü §
¡
Ó+ò 	$‰FˆB�Ø×Ñ˜r˜dÕ#ð	$ð 	×Ñœ4Ÿ;™;×=Ñ=¸lÌD×L\ÑL\×LbÑLbÐklÐjmÐ=ÓnÐoÔpØ×Ñä—‘×2Ñ2Ø'Ü×$Ñ$×*Ñ*ÚEð 3ó ðô	
ð 	�‰˜Ô Ø�
‰
×Ñ˜*Ô%âò3
Ð/ð ˆ	Ø—I‘Ió (	%ˆDØ�|‰|˜{Ó*Ü" 4›�ØŸ™›ò &�Að Ð.Ò.Ú$àÐ GÒGð Ð 0Ò0Ü"ŸN™NØ"6°q°cÐ9{Ð |ôð # 1™Ið!&ð$ �Ø—
‘
˜4Ÿ:™:Ô&ñ "Ü˜c›( Qš,ØŸ
™
 B 4Ô(ô ˜c›( Q›,ä˜3“x !’|ØŸ
™
 L >Ô2Ü˜3“x !’|ØŸ
™
Ð$7Ð#8Ô9ä—{‘{×,Ñ,Ø0ØØ—K‘Kñð Ÿ™ð	ð
 ñ�ð ×Ñ˜d˜VÖ$ðQ(	%ðR 	�‰˜ÔØ�	‰	×Ñ˜Ô#àr*   c                 ó  — t        «       }g }t        | j                  «      D ��ci c]  \  }}|j                  |“Œ }}}i }i }| j                  D ]N  }|j                  D ]$  }	|	sŒ|	|vr|g||	<   Œ||	   j                  |«       Œ& |j                  D ]
  }
|
sŒ|||
<   Œ ŒP | j                  D �]G  }|j                  dk(  sŒ|j                  d   r|j                  d   sŒ3|j                  d   |j                  d   }}d}d|v rO| j                  D ]?  }|j                  dk(  sŒ|j                  d   |k(  sŒ&|j                  d   j                  } n& n$| j                  D ]  }|j                  |k(  sŒ|} n |€ŒÍt        j                  j                  |«      }|j                  dk(  sŒü|j                  «       dv s�Œ|j                  d   |v s�Œ#||   }|j                  dk(  r|j                  d   s�ŒH|j                  d   |v r¯|j                  d   j!                  d	«      s|j                  d   j!                  d
«      rs|j                  «       dk(  r`|j#                  |j                  d   «       |j                  |«       t%        ||j                  d      «      dk(  r|j                  |«       �Œ|j                  d   |vr�Œ||j                  d      }|j                  dk(  r|j                  d   s�ŒM||j                  d      }|j                  dk(  r|j                  d   s�Œ|j                  d   |v s�Œ’|j                  d   j!                  d	«      s |j                  d   j!                  d
«      s�ŒÐ|j                  «       dk(  s�Œå|j#                  |j                  d   «       |j'                  |||g«       t%        ||j                  d      «      dk(  s�Œ7|j                  |«       �ŒJ ||fS c c}}w )az  Correct graph which originally use dim of past_seq_len from input_ids's shape which is fixed to max_seq_len after
       shared past/present buffer

    Args:
        subg (GraphProto): GraphProto of the decoder subgraph
    return:
        tensor_names_to_rename : set of tensor names which is equal to past_sequence_length
        nodes_to_remove : list of node to remove
    ÚGatherr\   r   NÚ	Constant_ÚConstant>   r\   r¶   ÚShaper  r  r¶   ÚReshaperµ   )Úsetrü   rÁ   r½   r+  rˆ   r›   r¿   r[  rT  r(  r«   r.  rÍ   ÚsizeÚitemÚ
startswithÚaddr‰   r‡   )rl  Útensor_names_to_renameÚnodes_to_removeÚindexÚinpÚgraph_input_namesÚinput_name_to_nodesr¾   r+  Ú
input_nameÚoutput_nameÚshape_tensor_nameÚshape_index_nameÚini_gather_indicesÚ
const_noderJ  Úgather_indices_arrÚ
shape_nodeÚreshape_nodeÚtranspose_nodes                       r(   Úfind_past_seq_len_usager«  ø  s  € ô !›UÐØ€Oä;DÀTÇZÁZÓ;P×Q©Z¨U°C˜Ÿ™ 5™ÐQÐÑQàÐØÐØ—	‘	ò 	8ˆØŸ*™*ò 	AˆJÚØÐ%8Ñ8Ø7;°fÐ'¨
Ò3à'¨
Ñ3×:Ñ:¸4Õ@ð	Að  Ÿ;™;ò 	8ˆKÚØ37Ð# KÒ0ñ	8ð	8ð —	‘	ó H!ˆð �<‰<˜8Ó#Ø—:‘:˜a’=¨¯
©
°1ªØð 48·:±:¸a±=À$Ç*Á*ÈQÁ-Ð/ÐØ!%ÐØÐ.Ñ.à"&§)¡)ò �JØ!×)Ñ)¨ZÓ7¸J×<MÑ<MÈaÑ<PÐTdÓ<dØ-7×-AÑ-AÀ!Ñ-D×-FÑ-FÐ*Ùñð #×.Ñ.ò �FØ—{‘{Ð&6Ó6Ø-3Ð*Ùðð "Ð)ØÜ!%×!2Ñ!2×!;Ñ!;Ð<NÓ!OÐð #×'Ñ'¨1Ó,Ø&×+Ñ+Ó-°Ó7Ø—J‘J˜q‘MÐ%8Ó8à0Ð1BÑC�
Ø"×*Ñ*¨gÒ5¸*×:JÑ:JÈ1Ò:MÙð ×$Ñ$ QÑ'Ð+<Ñ<à"×(Ñ(¨Ñ+×6Ñ6Ð7GÔHØ%×+Ñ+¨AÑ.×9Ñ9Ð:LÔMà*×/Ñ/Ó1°QÒ6ð +×.Ñ.¨t¯{©{¸1©~Ô>Ø#×*Ñ*¨4Ô0ÜÐ.¨z×/@Ñ/@ÀÑ/CÑDÓEÈÒJØ'×.Ñ.¨zÔ:Ù à×#Ñ# AÑ&Ð.AÑAÙØ2°:×3CÑ3CÀAÑ3FÑG�Ø$×,Ñ,°	Ò9¸l×>PÑ>PÐQRÒ>SÙØ!4°\×5GÑ5GÈÑ5JÑ!K�Ø&×.Ñ.°+Ò=À.×BVÑBVÐWXÒBYÙð #×(Ñ(¨Ñ+Ð/@Ó@à&×,Ñ,¨QÑ/×:Ñ:Ð;KÔLØ)×/Ñ/°Ñ2×=Ñ=Ð>PÖQà*×/Ñ/Ó1°QÔ6ð +×.Ñ.¨t¯{©{¸1©~Ô>Ø#×*Ñ*¨D°*¸lÐ+KÔLÜÐ.¨~×/DÑ/DÀQÑ/GÑHÓIÈQÔNØ'×.Ñ.¨~Ô>Ù ðQH!ðT " ?Ð2Ð2ùós Rs   ¥O;r°   Úpast_seq_len_namec                 ób  — d}t        t        d„ | j                  j                  j                  «      «      }|D ]„  }t        |j                  «      dk  r4|j                  j                  d«       t        |j                  «      dk  rŒ4|j                  j                  |«       |j                  j                  |«       Œ† | j                  j                  j                  j                  t        j                  j                  |t        j                  g d¢¬«      «       | j                  «        | S )Nr�  c                 ó    — | j                   dk(  S ©NÚMultiHeadAttention©r¿   ©r+  s    r(   ú<lambda>z.add_cache_indirection_to_mha.<locals>.<lambda>d  ó   € ¨¯©Ð9MÑ)M€ r*   r·   r   ©rƒ  r€  Úmax_sequence_lengthrq  )ÚlistÚfilterr°   r¼   r+  r‰   rÁ   rˆ   r«   r0  r1  r   rý   Útopological_sort)r°   r¬  Úcache_indirection_nameÚ	mha_nodesr+  s        r(   Úadd_cache_indirection_to_mhar¼  a  sí   € à0ÐÜ”VÑMÈuÏ{É{×O`ÑO`×OeÑOeÓfÓg€IØò 2ˆô �$—*‘*‹o Ò!Ø�J‰J×Ñ˜bÔ!ô �$—*‘*‹o Ó!à�
‰
×ÑÐ+Ô,Ø�
‰
×ÑÐ0Õ1ð2ð 
‡K�K×Ñ×Ñ×"Ñ"Ü�‰×*Ñ*Ø"¤K×$5Ñ$5Ò=pð 	+ó 	
ôð
 
×ÑÔØ€Lr*   r¹   Úskip_node_idxsc                 óx  — d}g }t        t        d„ | j                  j                  j                  «      «      }t        |«      D �]³  \  }}||v rŒd}|j                  D ]  }	|	j                  dk(  sŒ|	j                  } n |}
|
dk(  rO| j                  j                  j                  D ],  }|j                  |j                  d   k(  sŒ |j                  }
 n d}| j                  j                  j                  D ]W  }|j                  |j                  d   k(  sŒ |j                  j                  j                  j                  d   j                   } n t#        |j$                  «      dk  r4|j$                  j'                  d	«       t#        |j$                  «      dk  rŒ4|› d
|dz  › �}|j$                  j'                  |«       |j'                  t(        j*                  j-                  ||
d|d|g¬«      «       �Œ¶ | j                  j                  j$                  j/                  |«       | j1                  «        | S )NÚoutput_cross_qkc                 ó    — | j                   dk(  S r¯  r±  r²  s    r(   r³  z&add_output_qk_to_mha.<locals>.<lambda>z  r´  r*   r   r…  rì   Útarget_sequence_lengthr\   r¶   r   Ú_rƒ  Úsequence_lengthrq  )r·  r¸  r°   r¼   r+  rü   r[  r½   r  r(  rÁ   rÃ   r8   rÿ   r/  rf  rh  r‰   r›   rˆ   r«   r0  r1  r‡   r¹  )r°   r¹   r½  Úoutput_qk_basenameÚ
output_qksr»  Úidxr+  r…  ÚattÚoutput_qk_dtyper  rÁ  Úoutput_qk_names                 r(   Úadd_output_qk_to_mharÊ  v  s  € à*ÐØ€JÜ”VÑMÈuÏ{É{×O`ÑO`×OeÑOeÓfÓg€IÜ˜yÓ)ó (
‰	ˆˆTà�.Ñ Øð ˆ	Ø—>‘>ò 	ˆCØ�x‰x˜;Ó&ØŸE™E�	Ùð	ð  ˆØ˜aÒØ—[‘[×&Ñ&×2Ñ2ò �Ø—6‘6˜TŸZ™Z¨™]Ó*Ø&'§k¡k�OÙðð ":ÐØ—‘×"Ñ"×(Ñ(ò 	ˆAØ�v‰v˜Ÿ™ A™Ó&Ø)*¯©×);Ñ);×)AÑ)A×)EÑ)EÀaÑ)H×)RÑ)RÐ&Ùð	ô �$—+‘+Ó Ò"Ø�K‰K×Ñ˜rÔ"ô �$—+‘+Ó Ó"ð /Ð/¨q°¸±°
Ð;ˆØ�‰×Ñ˜>Ô*Ø×ÑÜ�K‰K×.Ñ.ØØØ# YÐ0AÐCYÐZð /ó ö	
ðE(
ðT 
‡K�K×Ñ×Ñ×#Ñ# JÔ/Ø	×ÑÔØ€Lr*   c                 óþ  ‡‡— d}d}d}t        t        d„ | j                  j                  j                  «      «      d   }| j                  |g d¢g d¢«      }| j                  |dd	gdd
g«      }|�|}n|�|}nt        j                  d«       y |d   }|j                  dk(  �rr|d   }	| j                  |	ddgddg«      Š‰€t        j                  d«       y | j                  |	g d¢g d¢«      }
|
€t        j                  d«       y |
d   }‰|
d
d  k7  rt        j                  d«       y t        t        ˆfd„| j                  j                  j                  «      «      d   }| j                  j                  j                  j                  |«       | j                  j                  j                  j                  ‰d   «       | j                  j                  j                  j                  ‰d
   «       ||	j                  d<   ||j                  d<   �n?| j                  |g d¢g d¢«      }|€t        j                  d«       y |d
   }| j                  |g d¢g d¢«      Š‰€t        j                  d«       y ‰d   }|dd  ‰d
d  k7  rt        j                  d«       y t        t        ˆfd„| j                  j                  j                  «      «      d   }| j                  j                  j                  j                  |«       t        t        ˆfd„| j                  j                  j                  «      «      d   }| j                  j                  j                  j                  |«       | j                  j                  j                  j                  ‰d
   «       | j                  j                  j                  j                  ‰d   «       | j                  j                  j                  j                  ‰d    «       | j                  j                  j                  j                  ‰d!   «       ||j                  d<   ||j                  d<   | j                  j                  j                  j                  t        j                  j                  |t        j                   d
g¬"«      «       t        j                  j#                  d#|g|g| j%                  d#«      ¬$«      }t        j                  j                  |t        j                   g ¬"«      }t        j                  j#                  d%|g|g| j%                  d%«      t        j&                  ¬&«      }t        j                  j                  |t        j&                  g ¬"«      }| j                  j                  j                  j)                  ||g«       | j                  j                  j*                  j)                  ||g«       | j-                  «        | |fS )'Nrp  Úpast_seq_len_int32Úpast_seq_len_int64c                 ó    — | j                   dk(  S )NÚLayerNormalizationr±  )Úns    r(   r³  z*fix_past_sequence_length.<locals>.<lambda>Ð  s   €  §¡Ð.BÑ!B€ r*   r   )ÚAddr‘  ÚTileÚExpandÚ	UnsqueezeÚRange)r   r\   r\   r   r   r   rÑ  ÚSlicer\   zBCannot identify base path for fixing past_sequence_length subgraphrb   rÕ  r‘  r”  zDCannot identify gather path for fixing past_sequence_length subgraph)rÑ  r‘  r”  ©r\   r   r   zACannot identify add path for fixing past_sequence_length subgraphz]Gather path and add path do not share the same nodes for calculating the past_sequence_lengthc                 óH   •— | j                   d   ‰d   j                  d   k(  S ©Nr   r\   ©r›   rÁ   )rÐ  Úgather_paths    €r(   r³  z*fix_past_sequence_length.<locals>.<lambda>  s$   ø€ °1·8±8¸A±;À+ÈaÁ.×BVÑBVÐWXÑBYÑ3Y€ r*   )rÔ  rÑ  r‘  r”  r•  rµ   )r¶   r   r   r   r   r   zGCannot identify input_ids path for fixing past_sequence_length subgraph)rÔ  r‘  r”  r•  rµ   )r\   r   r   r   r   zFCannot identify past_key path for fixing past_sequence_length subgraphr¶   ziThe input_ids path and past_key path do not share the same nodes for calculating the past_sequence_lengthc                 óH   •— | j                   d   ‰d   j                  d   k(  S rÙ  rÚ  ©rÐ  Úpast_key_paths    €r(   r³  z*fix_past_sequence_length.<locals>.<lambda>'  s&   ø€ °1·8±8¸A±;À-ÐPQÑBR×BXÑBXÐYZÑB[Ñ3[€ r*   c                 óH   •— | j                   d   ‰d   j                  d   k(  S )Nr   éþÿÿÿr\   rÚ  rÝ  s    €r(   r³  z*fix_past_sequence_length.<locals>.<lambda>)  s&   ø€ °A·H±H¸Q±KÀ=ÐQSÑCT×CZÑCZÐ[\ÑC]Ñ4]€ r*   rì   r^   rq  ÚSqueeze©ÚinputsÚoutputsr½   ÚCast©rã  rä  r½   Úto)r·  r¸  r°   r¼   r+  Úmatch_parent_pathr‹   rŒ   r¿   r,  rÁ   rˆ   r«   r0  r1  r   rý   rt  Úcreate_node_nameÚINT64r‡   r-  r¹  )r°   r¬  rÌ  rÍ  r+  Úbase_path_hfÚbase_path_oaiÚ	base_pathÚ	base_nodeÚ
range_nodeÚadd_pathÚadd_nodeÚconstant_in_gatherÚinput_ids_pathÚunsqueeze_nodeÚconstant_in_reshapeÚsqueeze_nodeÚsqueeze_outputÚ	cast_nodeÚcast_outputrÛ  rÞ  s                       @@r(   Úfix_past_sequence_lengthrú  ª  sW  ù€ ðD /ÐØ-ÐØ-Ðä”ÑBÀEÇKÁK×DUÑDU×DZÑDZÓ[Ó\Ð]^Ñ_€Dà×*Ñ*ØÚAÚó€Lð
 ×+Ñ+ØØ	�ÐØ	
ˆAˆó€Mð
 ÐØ ‰	Ø	Ð	"Ø!‰	ä�‰ÐXÔYØØ˜"‘€Ià×Ñ˜GÓ#à˜r‘]ˆ
à×-Ñ-ØØ�wÐØ�ˆFó
ˆð
 ÐÜ�K‰KÐ^Ô_Øà×*Ñ*ØÚ&Úó
ˆð
 ÐÜ�K‰KÐ[Ô\ØØ˜A‘;ˆà˜( 1 2˜,Ò&Ü�K‰KÐwÔxØô "¤&Ó)YÐ[`×[fÑ[f×[lÑ[l×[qÑ[qÓ"rÓsÐtuÑvÐØ�‰×Ñ×Ñ×%Ñ%Ð&8Ô9Ø�‰×Ñ×Ñ×%Ñ% k°!¡nÔ5Ø�‰×Ñ×Ñ×%Ñ% k°!¡nÔ5ð 1ˆ
×Ñ˜ÑØ.ˆ�‰�qÓð ×0Ñ0ØÚKÚó
ˆð
 Ð!Ü�K‰KÐaÔbØØ! !Ñ$ˆà×/Ñ/ØÚDÚó
ˆð
 Ð Ü�K‰KÐ`ÔaØØ& qÑ)ˆà˜!˜"Ð ¨q¨rÐ!2Ò2Ü�K‰KØ{ôð ô "¤&Ó)[Ð]b×]hÑ]h×]nÑ]n×]sÑ]sÓ"tÓuÐvwÑxÐØ�‰×Ñ×Ñ×%Ñ%Ð&8Ô9Ü"¤6Ó*]Ð_d×_jÑ_j×_pÑ_p×_uÑ_uÓ#vÓwØñ
Ðð 	�‰×Ñ×Ñ×%Ñ%Ð&9Ô:Ø�‰×Ñ×Ñ×%Ñ% m°AÑ&6Ô7Ø�‰×Ñ×Ñ×%Ñ% m°AÑ&6Ô7Ø�‰×Ñ×Ñ×%Ñ% m°AÑ&6Ô7Ø�‰×Ñ×Ñ×%Ñ% m°AÑ&6Ô7ð #5ˆ×Ñ˜QÑØ.ˆ�‰�qÑð 
‡K�K×Ñ×Ñ×"Ñ"Ü�‰×*Ñ*Ð+<¼k×>OÑ>OÐXYÐWZÐ*Ó[ôô
 —;‘;×(Ñ(ØØ!Ð"Ø#Ð$Ø×#Ñ# IÓ.ð	 )ó €Lô —[‘[×7Ñ7Ð8JÌK×L]ÑL]ÐegÐ7Óh€NÜ—‘×%Ñ%ØØ"Ð#Ø#Ð$Ø×#Ñ# FÓ+Ü×Ñð &ó €Iô —+‘+×4Ñ4Ð5GÌ×IZÑIZÐbdÐ4Óe€Kð 
‡K�K×Ñ×Ñ×!Ñ! <°Ð";Ô<Ø	‡K�K×Ñ× Ñ ×'Ñ'¨¸Ð(EÔFØ	×ÑÔØÐ#Ð#Ð#r*   c                 óž  — d}d}| j                   j                  j                  j                  t        j
                  j                  |t        j                  dg¬«      t        j
                  j                  |t        j                  g d¢¬«      g«       t        t        d„ | j                   j                  j                  «      «      }t        |«      D �]e  \  }}d}|j                  D ]  }|j                  dk(  sŒ|j                  } n d	|d
z  › �}	t        j
                  j                  |	t        j                   d|ddg¬«      }
|d
z  dk(  r/| j                   j                  j"                  j%                  |
«       t        j
                  j'                  d|j                  d   |j                  d   |j                  d
   ddt)        |j                  «      dkD  r|j                  d   ndt)        |j                  «      dkD  r|j                  d   nd||||j                  d   g|j"                  d   t)        |j                  «      dkD  r|j"                  d   ndt)        |j                  «      dkD  r|j"                  d
   nd|d
z  dk(  r|	ndg|j                  j+                  dd«      d||d
z  d¬«      }|d
z  dk(  r|j"                  j-                  d«       | j                   j                  j                  j-                  |«       | j                   j                  j                  j                  |g«       �Œh | j/                  «        | S )Nr€  r�  r\   rq  rµ  c                 ó    — | j                   dk(  S r¯  r±  r²  s    r(   r³  z(replace_mha_with_dmmha.<locals>.<lambda>b  r´  r*   r   r…  Úoutput_cross_qk_r¶   rƒ  zencode_sequence_length / 2ÚDecoderMaskedMultiHeadAttentionr   r^   rN  rO  rì   r°  úcom.microsoft)rã  rä  r½   r^  r…  Ú	output_qkrV   )r°   r¼   rÁ   r‡   r«   r0  r1  r   rý   r·  r¸  r+  rü   r[  r½   r  rþ   r›   rˆ   rt  r‰   Úreplacer,  r¹  )r°   r¬  r€  r�  r»  rÆ  r+  r…  rÇ  Úqk_output_nameÚ	qk_outputÚ
dmmha_nodes               r(   Úreplace_mha_with_dmmhar  S  sÜ  € à€JØ+Ðà	‡K�K×Ñ×Ñ×"Ñ"ä�K‰K×.Ñ.¨z¼;×;LÑ;LÐUVÐTWÐ.ÓXÜ�K‰K×.Ñ.Ø!¤;×#4Ñ#4Ò<oð /ó ð	
ôô ”VÑMÈuÏ{É{×O`ÑO`×OeÑOeÓfÓg€IÜ˜yÓ)ó 14‰	ˆˆTàˆ	Ø—>‘>ò 	ˆCØ�x‰x˜;Ó&ØŸE™E�	Ùð	ð ,¨C°1©H¨:Ð6ˆÜ—K‘K×6Ñ6ØœK×-Ñ-°lÀIÈqÐRnÐ5oð 7ó 
ˆ	ð �‰7�aŠ<Ø�K‰K×Ñ×$Ñ$×+Ñ+¨IÔ6ô —[‘[×*Ñ*Ø-à—
‘
˜1‘Ø—
‘
˜1‘Ø—
‘
˜1‘ØØÜ!$ T§Z¡Z£°1Ò!4�—
‘
˜1’¸"Ü!$ T§Z¡Z£°1Ò!4�—
‘
˜1’¸"Ø!ØØ!Ø—
‘
˜1‘ðð —‘˜A‘Ü"% d§j¡j£/°AÒ"5�—‘˜A’¸2Ü"% d§j¡j£/°AÒ"5�—‘˜A’¸2Ø"%¨¡'¨Q¢,‘°Bð	ð —‘×"Ñ"Ð#7Ð9ZÓ[Ø"ØØ˜Q‘wØ&'ð3 +ó 
ˆ
ð6 �‰7�aŠ<à×Ñ×$Ñ$ RÔ(à�‰×Ñ×Ñ×%Ñ% dÔ+Ø�‰×Ñ×Ñ×%Ñ% z lÖ3ðc14ðf 
×ÑÔØ€Lr*   Ú	attn_maskÚkv_num_headsÚ
world_sizeÚwindow_sizec                 ó†  — | j                  t        j                  j                  dt        j
                  dgdg¬«      «       t        j                  j                  d|dg|dz   g| j                  d«      ¬«      }t        j                  j                  d|dz   dgdg| j                  d«      ¬«      }t        j                  j                  d	dgd
g| j                  d	«      t        j                  ¬«      }t        j                  j                  d|g|dz   g| j                  d«      ¬«      }t        j                  j                  d|dz   dgdg| j                  d«      d¬«      }	t        j                  j                  d	dgdg| j                  d	«      t        j                  ¬«      }
| j                  j                  j                  j                  |||||	|
g«       t        t        d„ | j                  j                  j                  «      «      }t        |«      D �]   \  }}| j!                  |g d¢g d¢«      }| j!                  |ddgddg«      }d\  }}}|�|\  }}}n|�|\  }}| j!                  |g d¢g d¢«      }| j!                  |ddgddg«      }d\  }}}|�|\  }}}n|�|\  }}| j!                  |ddgddg«      }| j!                  |dgdg«      }d\  }}|�|\  }}n|�|d   }d}|�/|�-|j"                  D ]  }|j$                  dk(  sŒ|j&                  }Œ  d}|j"                  D ]  }|j$                  dk(  sŒ|j&                  }Œ  |j(                  d   |j(                  d   k(  xr |j(                  d   |j(                  d   k(  }|d uxr
 |d uxr |d u} |d u xr
 |d u xr |d u }!d\  }"}#}$|�rJ| s|!�rEt+        j,                  | j/                  |j(                  d   «      «      }%t+        j,                  | j/                  |j(                  d   «      «      }&t+        j,                  | j/                  |j(                  d   «      «      }'|%j0                  d    }(t3        j4                  |%|&|'fd¬!«      j7                  |(d"|(z  «      })t        j8                  j;                  |)d#|› �¬$«      })| j                  |)«       t        j                  j                  d|j(                  d   |)j$                  g|)j$                  › d%�g| j                  d«      ¬«      }*| j                  j                  j                  j                  |*g«       | j                  j                  j                  j=                  |«       | j                  j                  j                  j=                  |«       | j                  j                  j                  j=                  |«       |*j>                  d   }"| �rFt+        j,                  | j/                  |j(                  d   «      «      }+t+        j,                  | j/                  |j(                  d   «      «      },t+        j,                  | j/                  |j(                  d   «      «      }-|+j0                  d    }(t3        j4                  |+|,|-fd¬!«      j7                  d"|(z  «      }.t        j8                  j;                  |.d&|› �¬$«      }.| j                  |.«       t        j                  j                  d|*j>                  d   |.j$                  g|.j$                  › d%�g¬'«      }/| j                  j                  j                  j                  |/g«       | j                  j                  j                  j=                  |«       | j                  j                  j                  j=                  |«       | j                  j                  j                  j=                  |«       |/j>                  d   }"n-|j>                  d   }"|j>                  d   }#|j>                  d   }$t        j                  j                  d(|"|#|$|j(                  d)   |j(                  d*   |j>                  d   |
j>                  d   |�|j(                  d   nd+|�|j(                  d"   nd+g	|j>                  |j$                  jA                  d,d(«      d-||z  |dk(  r||z  n||z  |tC        |d uxr |d u«      |¬.«
      }0| j                  j                  j                  j=                  |«       | j                  j                  j                  j                  |0g«       |�/| j                  j                  j                  j=                  |«       |€�Œr| j                  j                  j                  j=                  |«       �Œ£ | S )/NÚoner\   ©r½   rÃ   rÅ   ÚvalsÚ	ReduceSumÚ	_row_sumsrâ  ÚSubÚseqlens_k_int64rå  Ú	seqlens_kræ  r”  Ú_shaper‘  Útotal_seq_len_int64r   )rã  rä  r½   r»   Útotal_seq_lenc                 ó    — | j                   dk(  S r¯  r±  r²  s    r(   r³  z&replace_mha_with_gqa.<locals>.<lambda>   r´  r*   )ÚRotaryEmbeddingrÑ  r´   )r   r   r   r  r´   )NNNr×  rÑ  r¶   ©NNÚinterleavedr…  )r   r   r   rb   rº   rì   ÚQKV_Weight_©r½   Ú_outputÚ	QKV_Bias_)rã  rä  ÚGroupQueryAttentionrN  rO  r   r°  rÿ  )	rã  rä  r½   r^  r…  r  Úlocal_window_sizeÚ	do_rotaryÚrotary_interleaved)"Úadd_initializerr«   r0  Úmake_tensorr   rê  rt  ré  rý   r°   r¼   r+  r‡   r·  r¸  rü   rè  r[  r½   r  rÁ   r   rÍ   rÀ   r/  rÉ   ÚstackÚreshaper.  Ú
from_arrayr,  r›   r  ru   )1r°   r  r  r  r	  Úreduce_sum_nodeÚsub_nodeÚseqlen_k_cast_noder¨  Úgather_nodeÚtotal_seqlen_cast_noder»  rÆ  r+  Úq_path_1Úq_path_2Úq_rotaryÚq_addÚq_matmulÚk_path_1Úk_path_2Úk_rotaryÚk_addÚk_matmulÚv_path_1Úv_path_2Úv_addÚv_matmulr  rÇ  r…  Úroot_input_is_sameÚall_paths_have_biasÚall_paths_have_no_biasÚq_input_to_attentionÚk_input_to_attentionÚv_input_to_attentionÚqwÚkwÚvwrf  Ú
qkv_weightÚpacked_matmul_nodeÚqbÚkbÚvbÚqkv_biasÚpacked_add_nodeÚgqa_nodes1                                                    r(   Úreplace_mha_with_gqarK  š  s\	  € ð& 
×ÑÜ�‰×ÑØÜ!×'Ñ'Ø�Ø�ð	 	 ó 	
ôô —k‘k×+Ñ+ØØ˜5Ð!Ø˜[Ñ(Ð)Ø×#Ñ# KÓ0ð	 ,ó €Oô �{‰{×$Ñ$ØØ˜KÑ'¨Ð/Ø"Ð#Ø×#Ñ# EÓ*ð	 %ó €Hô Ÿ™×.Ñ.ØØ!Ð"Ø�Ø×#Ñ# FÓ+Ü×Ñð /ó Ðô —‘×&Ñ&ØØˆ{Ø˜XÑ%Ð&Ø×#Ñ# GÓ,ð	 'ó €Jô —+‘+×'Ñ'ØØ˜HÑ$ eÐ,Ø&Ð'Ø×#Ñ# HÓ-Øð (ó €Kô "Ÿ[™[×2Ñ2ØØ%Ð&Ø Ð!Ø×#Ñ# FÓ+Ü×Ñð 3ó Ðð 
‡K�K×Ñ×Ñ×!Ñ!àØØØØØ"ð	
ô	ôH ”VÑMÈuÏ{É{×O`ÑO`×OeÑOeÓfÓg€IÜ˜yÓ)ó B4‰	ˆˆTà×*Ñ*¨4Ò1UÒW`ÓaˆØ×*Ñ*¨4Ð2CÀXÐ1NÐQRÐTUÐPVÓWˆà$4Ñ!ˆ�%˜ØÐØ(0Ñ%ˆH�e™XØÐ!Ø!)ÑˆH�hð ×*Ñ*¨4Ò1UÒW`ÓaˆØ×*Ñ*¨4Ð2CÀXÐ1NÐQRÐTUÐPVÓWˆà$4Ñ!ˆ�%˜ØÐØ(0Ñ%ˆH�e™XØÐ!Ø!)ÑˆH�hð ×*Ñ*¨4°%¸Ð1BÀQÈÀFÓKˆØ×*Ñ*¨4°(°¸a¸SÓAˆà$‰ˆˆxØÐØ&‰OˆE‘8ØÐ!Ø ‘{ˆHð ˆØÐ HÐ$8Ø×)Ñ)ò (�Ø—8‘8˜}Ó,Ø"%§%¡%‘Kð(ð
 ˆ	Ø—>‘>ò 	"ˆCØ�x‰x˜;Ó&ØŸE™E‘	ð	"ð
 &Ÿ^™^¨AÑ.°(·.±.ÀÑ2CÑCÒnÈÏÉÐWXÑHYÐ]e×]kÑ]kÐlmÑ]nÑHnÐð $¨4Ð/Ò[°EÀÐ4EÒ[È%ÐW[ÐJ[ÐØ!&¨$ Ò!R°5¸D°=Ò!RÀUÈdÀ]Ðð LVÑHÐÐ2Ð4HÚÑ#6Ò:PÜ×%Ñ% e×&;Ñ&;¸H¿N¹NÈ1Ñ<MÓ&NÓOˆBÜ×%Ñ% e×&;Ñ&;¸H¿N¹NÈ1Ñ<MÓ&NÓOˆBÜ×%Ñ% e×&;Ñ&;¸H¿N¹NÈ1Ñ<MÓ&NÓOˆBà—(‘(˜2‘,ˆCÜŸ™ 2 r¨2 ,°QÔ7×?Ñ?ÀÀQÈÁWÓMˆJÜ×*Ñ*×5Ñ5°jÈÐUXÐTYÐGZÐ5Ó[ˆJØ×!Ñ! *Ô-ä!%§¡×!6Ñ!6ØØ Ÿ™ qÑ)¨:¯?©?Ð;Ø&ŸO™OÐ,¨GÐ4Ð5Ø×+Ñ+¨HÓ5ð	 "7ó "Ðð �K‰K×Ñ×"Ñ"×)Ñ)Ð+=Ð*>Ô?Ø�K‰K×Ñ×"Ñ"×)Ñ)¨(Ô3Ø�K‰K×Ñ×"Ñ"×)Ñ)¨(Ô3Ø�K‰K×Ñ×"Ñ"×)Ñ)¨(Ô3Ø#5×#<Ñ#<¸QÑ#?Ð ò #Ü ×)Ñ)¨%×*?Ñ*?ÀÇÁÈAÁÓ*OÓP�Ü ×)Ñ)¨%×*?Ñ*?ÀÇÁÈAÁÓ*OÓP�Ü ×)Ñ)¨%×*?Ñ*?ÀÇÁÈAÁÓ*OÓP�à—h‘h˜r‘l�ÜŸ8™8 R¨¨R L°qÔ9×AÑAÀ!ÀcÁ'ÓJ�Ü×,Ñ,×7Ñ7¸ÈÐSVÐRWÐGXÐ7ÓY�Ø×%Ñ% hÔ/Ü"&§+¡+×"7Ñ"7ØØ.×5Ñ5°aÑ8¸(¿-¹-ÐHØ (§¡˜¨gÐ6Ð7ð #8ó #�ð
 —‘×!Ñ!×&Ñ&×-Ñ-¨Ð.?Ô@Ø—‘×!Ñ!×&Ñ&×-Ñ-¨eÔ4Ø—‘×!Ñ!×&Ñ&×-Ñ-¨eÔ4Ø—‘×!Ñ!×&Ñ&×-Ñ-¨eÔ4Ø'6×'=Ñ'=¸aÑ'@Ñ$ð $,§?¡?°1Ñ#5Ð Ø#+§?¡?°1Ñ#5Ð Ø#+§?¡?°1Ñ#5Ð ô —;‘;×(Ñ(Ø!à$Ø$Ø$Ø—
‘
˜1‘Ø—
‘
˜1‘Ø"×)Ñ)¨!Ñ,Ø&×-Ñ-¨aÑ0Ø&.Ð&:�—‘ Ò"ÀØ&.Ð&:�—‘ Ò"Àð
ð —K‘KØ—‘×"Ñ"Ð#7Ð9NÓOØ"Ø :Ñ-Ø5AÀQÒ5F˜) zÒ1ÈLÐ\fÑLfØ)Ü˜(¨$Ð.ÒG°8À4Ð3GÓHØ*ð) )ó 
ˆð, 	�‰×Ñ×Ñ×%Ñ% dÔ+Ø�‰×Ñ×Ñ×%Ñ% x jÔ1àÐØ�K‰K×Ñ×"Ñ"×)Ñ)¨(Ô3ØÒØ�K‰K×Ñ×"Ñ"×)Ñ)¨(Ö3ðEB4ðH €Lr*   c           	      óÌ  — d}| j                   D �cg c]  }|j                  ‘Œ }}|dk  r3||   j                  d«      s|dz  }|dk  r||   j                  d«      sŒd}t        | j                  «      |z
  dz  }d|z  |z   }t        |«      D �ci c]"  }| j                   |dz  |z      j                  |“Œ$ }}t        d|› �«       t        | j                   |   «      }	t        d|	› �«       |	d   }
|	d   }|	d   }d}| j                  D �]7  }|j                  dk(  sŒ|j                   d   |v sŒ&t        d	|j                  › d
|j                  › �«       |dz  }||j                   d      }d|› �}dgdt        |j                  «      z
  z  }|j                  |«       |j                  j                  |«       |j                  j                  t        j                  j                  dd«      g«       t        j                  j!                  |t"        j$                  |
|d|g«      }| j                  j                  |g«       �Œ: ||k7  rt'        d|› d|› �«      ‚y c c}w c c}w )Nr\   rì   Úpastr¶   z    -- past_key_cross_inputs = zpast_key_cross_0_shape is r   rþ  z'    -- add cross QK output from: node: z with output: rý  r   r   z#Did not add cross QK for all layersz vs )rÁ   r½   r™  r‰   r›   rú   Úprintrk  r+  r¿   rˆ   r‡   r[  r«   r0  Úmake_attributer1  r   rþ   rû   )rl  Úinput_self_past_0ÚgirŸ  Úoutput_self_present_0Ú
num_layersÚinput_cross_past_0ÚlayerÚpast_key_cross_inputsÚinput_past_key_cross_0_shapeÚbatch_size_dimÚnum_heads_dimÚcross_seq_len_dimÚnum_layer_output_qkr+  Úcross_attention_out_nameÚappended_namesÚcross_attentions                     r(   Ú.update_decoder_subgraph_output_cross_attentionr_  ˆ  sˆ  € ØÐà+/¯:©:Ö6 R˜Ÿ›Ð6ÐÐ6Ø
˜aÒ
Ð(9Ð:KÑ(L×(WÑ(WÐX^Ô(_Ø˜QÑÐð ˜aÒ
Ð(9Ð:KÑ(L×(WÑ(WÐX^Õ(_àÐä�d—k‘kÓ"Ð%:Ñ:¸qÑ@€JØ˜Z™Ð*;Ñ;ÐÜafÐgqÓarÖsÐX]˜TŸZ™Z¨°©	Ð4FÑ(FÑG×LÑLÈeÑSÐsÐÐsÜ	Ð+Ð,AÐ+BÐ
CÔDä#+¨D¯J©JÐ7IÑ,JÓ#KÐ Ü	Ð&Ð'CÐ&DÐ
EÔFØ1°!Ñ4€NØ0°Ñ3€MØ4°QÑ7ÐàÐØ—	‘	ó 2ˆØ�L‰LÐ=Ó=ÀDÇJÁJÈqÁMÐUjÒDjÜÐ;¸D¿I¹I¸;ÀnÐUY×U`ÑU`ÐTaÐbÔcØ 1Ñ$ÐØ)¨$¯*©*°Q©-Ñ8ˆEØ)9¸%¸Ð'AÐ$Ø ˜T Q¬¨T¯[©[Ó)9Ñ%9Ñ:ˆNØ×!Ñ!Ð":Ô;Ø�K‰K×Ñ˜~Ô.Ø�N‰N×!Ñ!¤4§;¡;×#=Ñ#=¸kÈ1Ó#MÐ"NÔOä"Ÿk™k×@Ñ@Ø(Ü×!Ñ!Ø °Ð3DÐEóˆOð
 �K‰K×Ñ Ð0Ö1ð!2ð" ˜jÒ(ÜÐ>¸z¸lÈ$ÐObÐNcÐdÓeÐeð )ùòE 7ùò ts   ‘IÂ'I!c           
      ó  — d}| j                   D �cg c]  }|j                  ‘Œ }}|dk  r3||   j                  d«      s|dz  }|dk  r||   j                  d«      sŒd}t        t	        | j                   «      |z
  dz  «      }d|z  |z   }g }g }| j
                  D ]$  }	|	j                  dk(  sŒ|j                  |	g«       Œ& t	        |«      |k  ryd }
| j
                  D ]  }	|	j                  dk(  sŒ|	}
 n g d	¢}d
}t        | «      \  }}t	        |«      dkD  r±|D ]  }t        d|› d|› d�«       Œ |D ]'  }t        d|j                  › d|j                  › �«       Œ) t        j                  j                  ddgdgd¬«      }t        j                  j                  ddg|gdt        j                  ¬«      }|j                  ||g«       | j
                  D �]Û  }	t	        |	j                  «      dkD  rƒ|
��|	j                  d   |
j                   d   k(  rbt        j                  j                  ddgdgdt        j                  ¬«      }|j                  d   |	j                   d<   |j                  |g«       |	j                  dk(  �rät!        |	«      }|j#                  «       D ]
  }||vsŒ||= Œ |	j                   d   |	j                   d   |	j                   d   g}|j                  t	        |	j                   «      dkD  r|	j                   d   ndg«       |j                  t	        |	j                   «      dkD  r|	j                   d   ndg«       |j                  t	        |	j                   «      dkD  r|	j                   d   ndg«       |j                  t	        |	j                   «      dkD  r|	j                   d   ndg«       |j                  dg«       |j                  dg«       |j                  d g«       |j                  t	        |	j                   «      dkD  r|	j                   d   ndg«       d|d!<   t        j                  j                  d"||	j                  fd#|	j                  i|¤Ž}	|	|vs�Œ™t%        |	j                   «      D ]  \  }}||v sŒ||	j                   |<   Œ |j                  |	g«       �ŒÞ | j'                  d$«       | j
                  j                  |«       | j                   D �cg c]  }|j                  ‘Œ }}g }t%        | j                   «      D ]ƒ  \  }}||k\  rg||k  rbt)        |«      }t        j                  j+                  |j                  |j,                  j.                  j0                  |d   |d   d%|d   g¬&«      }|j                  |g«       Œ… d|vrK|j                  t        j                  j+                  dt        j                  j2                  dg¬'«      g«       d|vrK|j                  t        j                  j+                  dt        j                  j2                  dg¬'«      g«       d |vrL|j                  t        j                  j+                  d t        j                  j2                  g d(¢¬'«      g«       | j'                  d)«       | j                   j                  |«       g }t%        | j                  «      D ]~  \  }}||k\  rbt)        |«      }t        j                  j+                  |j                  |j,                  j.                  j0                  |d   |d   d%|d   g¬&«      }|j                  |g«       Œ€ | j'                  d*«       | j                  j                  |«       y+c c}w c c}w ),Nr\   rì   rM  r^   r¶   r°  FÚRelativePositionBiasr„  Ú#past_sequence_length_squeezed_int64r   zFound tensor name `z` to be renamed to `ú`zFound node to remove: type = z	, name = rá  rp  Úpast_sequence_length_squeezedÚ!node_past_sequence_length_squeezer  rå  Ú&node_past_sequence_length_squeeze_cast)r½   rç  Úpast_sequence_length_int64Úpast_sequence_length_castr   rM  rN  rO  r€  r�  rV   rþ  r½   r+  rn  ro  rq  r‚  rÁ   r›   T)rÁ   r½   r™  ru   r‰   r+  r¿   r‡   r«  rN  r«   r0  rt  r   rê  r›   rc  r‹  rü   rs  rk  r1  r8   rÿ   r   rý   )rl  rP  rQ  rŸ  Úoutput_self_past_0rS  rT  ry  Ú	old_nodesr+  Úrel_pos_bias_noder�  Útarget_squeezed_past_seq_namer›  rœ  Úname_to_renameÚnrrö  rø  r`  rŽ  r{  r�  r½   rž  Úorig_input_namesrw  r  ri  r/  rx  s                                  r(   Ú?update_decoder_subgraph_share_buffer_and_use_decoder_masked_mharp  ±  s:  € ØÐà+/¯:©:Ö6 R˜Ÿ›Ð6ÐÐ6Ø
˜aÒ
Ð(9Ð:KÑ(L×(WÑ(WÐX^Ô(_Ø˜QÑÐð ˜aÒ
Ð(9Ð:KÑ(L×(WÑ(WÐX^Õ(_àÐä”c˜$Ÿ*™*“oÐ(9Ñ9¸QÑ>Ó?€JØ˜Z™Ð*;Ñ;Ðà€IØ€IØ—	‘	ò %ˆØ�<‰<Ð/Ó/Ø×Ñ˜d˜VÕ$ð%ô
 ˆ9ƒ~˜
Ò"Øð ÐØ—	‘	ò ˆØ�<‰<Ð1Ó1Ø $ÐÙðò
/Ð+ð %JÐ!Ü.EÀdÓ.KÑ+Ð˜OÜ
Ð!Ó" QÒ&Ø4ò 	nˆNÜÐ'¨Ð'7Ð7KÐLiÐKjÐjkÐlÕmð	nà!ò 	RˆBÜÐ1°"·*±*°¸YÀrÇwÁwÀiÐPÕQð	Rô —{‘{×,Ñ,ØØ#Ð$Ø,Ð-Ø4ð	 -ó 
ˆô —K‘K×)Ñ)ØØ,Ð-Ø*Ð+Ø9Ü× Ñ ð *ó 
ˆ	ð 	×Ñ˜,¨	Ð2Ô3à—	‘	ó 0%ˆÜˆt�{‰{Ó˜aÒÐ$5Ð$AÀdÇkÁkÐRSÁnÐXi×XoÑXoÐpqÑXrÒFrÜŸ™×-Ñ-ØØ'Ð(Ø-Ð.Ø0Ü×$Ñ$ð .ó ˆIð &×,Ñ,¨QÑ/ˆD�J‰J�q‰MØ×Ñ˜i˜[Ô)à�<‰<Ð/Ó/Ü˜t“_ˆFØ—[‘[“]ò "�ØÐCÒCØ˜q™	ð"ð —
‘
˜1‘Ø—
‘
˜1‘Ø—
‘
˜1‘ðˆCð �J‰J¬¨T¯Z©Z«¸1Ò)<˜Ÿ
™
 1šÀ"ÐEÔFØ�J‰J¬¨T¯Z©Z«¸1Ò)<˜Ÿ
™
 1šÀ"ÐEÔFØ�J‰J¬¨T¯Z©Z«¸1Ò)<˜Ÿ
™
 1šÀ"ÐEÔFØ�J‰J¬¨T¯Z©Z«¸1Ò)<˜Ÿ
™
 1šÀ"ÐEÔFØ�J‰JÐ.Ð/Ô0Ø�J‰J˜�~Ô&Ø�J‰JÐ+Ð,Ô-Ø�J‰J¬¨T¯Z©Z«¸1Ò)<˜Ÿ
™
 1šÀ"ÐEÔFà23ˆFÐ.Ñ/ä—;‘;×(Ñ(Ø1ØØ—‘ñð —Y‘Yð	ð
 ñˆDð �Ó&Ü(¨¯©Ó4ò F‘��tØÐ1Ò1Ø(E�D—J‘J˜uÒ%ðFð ×Ñ˜d˜VÖ$ða0%ðd 	‡O�O�FÔØ‡I�I×Ñ�YÔØ,0¯J©JÖ7 S˜Ÿ›Ð7ÐÐ7à€JÜ˜4Ÿ:™:Ó&ò  ‰ˆˆ2ØÐ!Ò! aÐ*<Ò&<Ü˜R“LˆEÜ—‘×3Ñ3Ø—‘ØŸ'™'×-Ñ-×7Ñ7Ø˜Q‘x  q¡¨=¸%À¹(ÐCð 4ó ˆBð
 	×Ñ˜2˜$Õð ð Ð%5Ñ5Ø×ÑÜ�[‰[×/Ñ/Ð0FÌ×HXÑHX×H^ÑH^ÐghÐfiÐ/ÓjÐkô	
ð Ð+Ñ+Ø×Ñœ4Ÿ;™;×=Ñ=¸lÌD×L\ÑL\×LbÑLbÐklÐjmÐ=ÓnÐoÔpØÐ"2Ñ2Ø×Ñä—‘×2Ñ2Ø'Ü×$Ñ$×*Ñ*ÚEð 3ó ðô	
ð 	‡O�O�GÔØ‡J�J×Ñ�jÔ!à€KÜ˜4Ÿ;™;Ó'ò !‰ˆˆ2ØÐ"Ò"Ü˜R“LˆEÜ—‘×3Ñ3Ø—‘ØŸ'™'×-Ñ-×7Ñ7Ø˜Q‘x  q¡¨=¸%À¹(ÐCð 4ó ˆBð
 	×Ñ˜B˜4Õ ð!ð 	‡O�O�HÔØ‡K�K×Ñ�{Ô#àùòs 7ùòZ 8s   ‘^Ó0^Úmodel_protoc                 óö  — t        | «      }|j                  «       }g }g }|j                  «       D �]„  }|j                  dk(  sŒd|j                  d   v rd|j                  d   v rŒ7||j                  d      }||j                  d      }||j                  d      }|j                  |j                  d   «      }	|j                  |j                  d   «      }
|j                  |j                  d   «      }|	r|
r|s yt        j                  |	«      }t        j                  |
«      }t        j                  |«      }t        j                  |||gd¬«      }|j                  d	d
¬«      }t        j                  j                  |dz   |	j                  dk(  rt        j                   nt        j"                  |j$                  d   |j$                  d   g|j'                  «       j)                  «       ¬«      }| j*                  j,                  j/                  |g«       t        j                  j1                  d	|j                  d   |dz   g|dz   g|¬«      }|j2                  d   |j                  d<   d|j                  d<   d|j                  d<   |j/                  |g«       |j/                  |||g«       �Œ‡ |j5                  |«       |j7                  |«       |j9                  «        |j;                  «        y)Nrþ  Úpast_key_crossr\   Úpast_value_crossr¶   r   Frº   r´   Ú
MatMul_QKV)Úname_prefixÚ_weightr  Ú_outrâ  r   T)r   r¾   Únodesr¿   rÁ   rÀ   r   rÍ   rÉ   rÌ   ré  r«   r0  r#  rÃ   r   rþ   rt   r/  ÚflattenÚtolistr¼   r(  r‡   rt  r›   Ú	add_nodesÚremove_nodesÚupdate_graphr¹  )rq  Ú
onnx_modelr¾   Únodes_to_addrœ  r+  r0  r5  r9  Úq_weightÚk_weightÚv_weightr@  rA  rB  rC  Úmatmul_node_nameÚweightrÒ   s                      r(   Úpack_qkv_for_decoder_masked_mhar†  P  s½  € Ü˜;Ó'€JØ$×8Ñ8Ó:Ðà€LØ€OØ× Ñ Ó"ó *CˆØ�<‰<Ð<Ó<Ø 4§:¡:¨a¡=Ñ0Ð5GÈ4Ï:É:ÐVWÉ=Ñ5XØØ*¨4¯:©:°a©=Ñ9ˆHØ*¨4¯:©:°a©=Ñ9ˆHØ*¨4¯:©:°a©=Ñ9ˆHà!×1Ñ1°(·.±.ÀÑ2CÓDˆHØ!×1Ñ1°(·.±.ÀÑ2CÓDˆHØ!×1Ñ1°(·.±.ÀÑ2CÓDˆHÙ¡©hÙä×%Ñ% hÓ/ˆBÜ×%Ñ% hÓ/ˆBÜ×%Ñ% hÓ/ˆBäŸ™¨¨R°¨¸1Ô=ˆJà)×:Ñ:¸8ÐQ]Ð:Ó^ÐÜ—[‘[×,Ñ,Ø%¨	Ñ1Ø08×0BÑ0BÀaÒ0Gœ;×,Ò,Ì[×M`ÑM`Ø ×&Ñ& qÑ)¨:×+;Ñ+;¸AÑ+>Ð?Ø×'Ñ'Ó)×0Ñ0Ó2ð	 -ó ˆFð ×Ñ×)Ñ)×0Ñ0°&°Ô:äŸ+™+×/Ñ/ØØ Ÿ™ qÑ)Ð+;¸iÑ+GÐHØ)¨FÑ2Ð3Ø%ð	 0ó ˆKð (×.Ñ.¨qÑ1ˆD�J‰J�q‰MØˆD�J‰J�q‰MØˆD�J‰J�q‰Mà×Ñ  Ô.Ø×"Ñ" H¨h¸Ð#AÖBðU*CðX ×Ñ˜Ô&Ø×Ñ˜OÔ,Ø×ÑÔà×ÑÔ!àr*   Údecoder_onnx_pathc                 ó.  — t        j                  | d¬«      }t        t        |j                  j
                  «      «      D ]»  }|j                  j
                  |   j                  dk(  s'|j                  j
                  |   j                  dk(  sŒP|j                  j
                  |   j                  j                  j                  j                  d   }|j                  d«      r|j                  «        d|_        Œ½ t        j                  || |¬«       y)aQ  Update the input shapes for the inputs "input_ids" and "position_ids" and make the sequence length dim value 1 for each of them.
       The decoder model will be over-written.

    Args:
        decoder_onnx_path (str): Path of GPT-2 decoder onnx model
        use_external_data_format(bool): output tensors to external data or not.
    Tr¤   rí   rî   r\   rg  r¨   )r«   r¬   rú   r‰   r¼   rÁ   r½   r8   rÿ   r/  rf  ÚHasFieldÚClearrh  r   r®   )r‡  rL   rÏ   r  Úshape_dim_protos        r(   Ú*update_input_shapes_for_gpt2_decoder_modelrŒ  ‹  sñ   € ô Ÿ/™/Ð*;ÐPTÔUÐÜ”3Ð*×0Ñ0×6Ñ6Ó7Ó8ò *ˆà×%Ñ%×+Ñ+¨AÑ.×3Ñ3°{ÒBØ"×(Ñ(×.Ñ.¨qÑ1×6Ñ6¸.ÓHà1×7Ñ7×=Ñ=¸aÑ@×EÑE×QÑQ×WÑW×[Ñ[Ð\]Ñ^ˆOð ×'Ñ'¨Ô4Ø×%Ñ%Ô'ð )*ˆOÕ%ð*ô ‡N�NØØØ6õð
 r*   Úinit_decoder_onnx_pathc           	      ó¸  — t        j                  | d¬«      }|j                  j                  d   j                  }t        |«      }|j                  «       }||v sJ ‚||   }|j                  dk7  ry|j                  |g d¢g d¢«      }|€|j                  |g d¢g d	¢«      }|€0|j                  |g d
¢g d¢«      }|€|j                  |g d¢g d¢«      }|€y|d   }	|	j                  dk(  }
|
sqd}|j                  |	g d¢|dddg«      }|€d}|j                  |	g d¢|dddg«      }|€d}|j                  |	g d¢|ddg«      }|€‡d}|j                  |	g d¢|ddg«      }nld}|j                  |	g d¢|ddg«      }|€d}|j                  |	g d¢|ddg«      }|€d}|j                  |	ddg|dg«      }|€d}|j                  |	ddg|dg«      }|€y|dk(  rdnd}|
s|j                  |	d|«      }n|j                  |	d|«      }|€y|d   }|d   }t         j                  j                  dt        j                  dgdg¬«      }t         j                  j                  dt        j                  dgdg¬«      }t         j                  j                  dt        j                  dgdg¬«      }t         j                  j                  dt        j                  dgdg¬«      }|j                  |«       |j                  |«       |j                  |«       |j                  |«       d|j                  d   z   }t         j                  j                  d|j                  d   ddddg|g|j!                  dd«      ¬«      }|
s|j                  d   n|j                  d    }d|j                  d   z   }t         j                  j                  d|ddddg|g|j!                  dd!«      ¬«      }|j#                  |«       |j#                  |«       |j%                  ||j                  d   |«       |j%                  |	||«       |j'                  «        t        j(                  |||¬"«       y)#a„  Generates the initial decoder GPT2 subgraph and saves it for downstream use.
       The initial decoder model will be saved to init_decoder_onnx_path.

    Args:
        decoder_onnx_path (str): Path of GPT-2 decoder onnx model
        init_decoder_onnx_path (str): Path of GPT-2 init decoder onnx model
        use_external_data_format(bool): output tensors to external data or not.
    Tr¤   r   r´   F)rå  rÏ  rÑ  rÑ  rå  r´   rå  ÚFastGelurå  r´   rå  rÏ  rÑ  )r   r   r   r\   r   r   r   r   r   r   r   r   r   )
rå  ÚSkipLayerNormalizationrå  r´   rå  r�  rå  r´   rå  r�  )
r   r   r\   r   r   r   r   r   r   r   )rÏ  rÑ  rÑ  r´   r�  r´   rÏ  rÑ  )r   r   r\   r   r   r   r   r   )r�  r´   r�  r´   r�  )r   r\   r   r   r   rb   r�  )rÑ  rå  r´   rr  r\   )rÑ  r´   rr  )rå  r´   rr  rr  rÑ  rà  ÚSliceLastTokenStartsr  ÚSliceLastTokenEndsÚSliceLastTokenAxesÚSliceLastTokenStepsÚedge_modified_rÖ  ÚGatherLastToken_0_râ  rì   ÚGatherLastToken_1_r¨   )r«   r¬   r¼   r›   r½   r   r¾   r¿   rè  rÂ   r0  r#  r   rý   r"  rt  ré  rñ  Úreplace_node_inputr¹  r®   )r‡  r�  rL   Úinit_decoder_model_protorÐ   Úgpt2_init_decoder_modelr¾   Úlogits_matmul_nodeÚ"logits_matmul_to_residual_add_pathÚresidual_add_nodeÚis_skiplayernorm_pathÚ&residual_add_to_attention_parent_indexÚresidual_add_to_attention_pathÚ residual_add_to_add_parent_indexÚadd_before_residual_addÚ	attentionÚmatmul_after_attentionÚslice_startsÚ
slice_endsÚ
slice_axesÚslice_stepsÚslice_0_output_nameÚslice_node_0Úadd_before_residual_add_outputÚslice_1_output_nameÚslice_node_1s                             r(   Úgenerate_gpt2_init_decoderr®  «  sš  € ô  $Ÿ™Ð/@ÐUYÔZÐà1×7Ñ7×>Ñ>¸qÑA×FÑFÐä'Ð(@ÓAÐà1×EÑEÓGÐØÐ!4Ñ4Ð4Ð4à,Ð-?Ñ@Ðð ×!Ñ! XÒ-Øð *A×)RÑ)RØò	
ò 	0ó#*Ð&ð* *Ð1Ø-D×-VÑ-VØòò +ó.
Ð*ð$ *Ð1à-D×-VÑ-VØò	ò %ó.
Ð*ð  .Ð5Ø1H×1ZÑ1ZØ"òò  ó
2Ð.ð *Ð1Øà:¸2Ñ>Ðð .×5Ñ5Ð9QÑQÐñ !Ø12Ð.Ø)@×)RÑ)RØÚ2Ø3°Q¸¸1Ð=ó*
Ð&ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ú6Ø7¸¸A¸qÐAó.Ð*ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ú.Ø7¸¸AÐ>ó.Ð*ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ú.Ø7¸¸AÐ>ó.Ñ*ð 23Ð.Ø)@×)RÑ)RØÚ+Ø3°Q¸Ð:ó*
Ð&ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ú/Ø7¸¸AÐ>ó.Ð*ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ø˜;Ð'Ø7¸Ð;ó.Ð*ð *Ð1Ø56Ð2Ø-D×-VÑ-VØ!Ø˜;Ð'Ø7¸Ð;ó.Ð*ð &Ð-Øà,RÐVWÒ,W¡qÐ]^Ð$ñ !Ø"9×"FÑ"FØ˜uÐ&Fó#
Ñð #:×"FÑ"FØØ$Ø,ó#
Ðð Ð&Øà.¨rÑ2€IØ;¸BÑ?Ðä—;‘;×*Ñ*Ø#Ü×#Ñ#ØˆSØˆTð	 +ó €Lô —‘×(Ñ(Ø!Ü×#Ñ#ØˆSØˆTð	 )ó €Jô —‘×(Ñ(Ø!Ü×#Ñ#ØˆSØˆSð	 )ó €Jô —+‘+×)Ñ)Ø"Ü×#Ñ#ØˆSØˆTð	 *ó €Kð ×+Ñ+¨LÔ9Ø×+Ñ+¨JÔ7Ø×+Ñ+¨JÔ7Ø×+Ñ+¨KÔ8ð +¨Y×-=Ñ-=¸aÑ-@Ñ@ÐÜ—;‘;×(Ñ(Øà×Ñ˜QÑØ"Ø Ø Ø!ð
ð %Ð%Ø$×5Ñ5°gÐ?SÓTð )ó €Lñ" 2GÐ×&Ñ& qÒ)ÐLc×LjÑLjÐklÑLmð #ð +Ð-D×-KÑ-KÈAÑ-NÑNÐÜ—;‘;×(Ñ(Øà*Ø"Ø Ø Ø!ð
ð %Ð%Ø$×5Ñ5°gÐ?SÓTð )ó €Lð ×$Ñ$ \Ô2Ø×$Ñ$ \Ô2ð ×.Ñ.Ð/EÀy×GWÑGWÐXYÑGZÐ\oÔpØ×.Ñ.Ð/@ÐB`ÐbuÔvð ×,Ñ,Ô.ô ‡N�NØ ØØ6õð
 r*   c                 ó  — t        d«      }t        |j                  «      }t        |j                  «      }t        |j                  «      }| j                  j
                  D ]‚  }|j                  j                  j                  j                  D ]S  }|j                  d«      sŒ|j                  ||||fv sŒ(t        |j                  «      }|j                  «        ||_        ŒU Œ„ | j                  j                  D ]‚  }|j                  j                  j                  j                  D ]S  }|j                  d«      sŒ|j                  ||||fv sŒ(t        |j                  «      }|j                  «        ||_        ŒU Œ„ y)zoMake dim_proto numeric.

    Args:
        model: T5 encoder and decoder model.
        config: T5 config.
    r\   rg  N)rn   r…  Úd_modelÚd_kvr¼   r›   r8   rÿ   r/  rf  r‰  rg  ru   rŠ  rh  rÁ   )	r°   ÚconfigrÃ  r…  Úhidden_sizeÚ	head_sizerJ  Ú	dim_protorh  s	            r(   Úmake_dim_proto_numeric_t5r¶  Ö	  sY  € ô ˜!“f€OÜ�F×$Ñ$Ó%€IÜ�f—n‘nÓ%€KÜ�F—K‘KÓ €Ià—+‘+×$Ñ$ò 
0ˆØŸ™×0Ñ0×6Ñ6×:Ñ:ò 		0ˆIØ×!Ñ! +Õ.°9×3FÑ3FØØØØð	Kò 4ô   	× 3Ñ 3Ó4�	Ø—‘Ô!Ø&/�	Õ#ñ		0ð
0ð —+‘+×#Ñ#ò 
0ˆØŸ™×0Ñ0×6Ñ6×:Ñ:ò 		0ˆIØ×!Ñ! +Õ.°9×3FÑ3FØØØØð	Kò 4ô   	× 3Ñ 3Ó4�	Ø—‘Ô!Ø&/�	Õ#ñ		0ñ
0r*   Úgeneration_typec                 ó,  — | j                   dk(  }|t        j                  k(  }|t        j                  k(  }|t        j                  k(  }| j
                  }t        j                  d|› �«       t        | j                  «      dk(  rƒ| j                  d   dk(  rq|rh| j                  t        j                  j                  k(  rAg d¢| _	        t        j                  d| j                  › �«       t        j                  d«       ng | _	        |s|r;|st        d	«      ‚| j                  rt        d
«      ‚| j                   rt        d«      ‚|r|r| j"                  st%        d«      ‚| j"                  r|st%        d«      ‚| j"                  r| j&                  st%        d«      ‚|rù| j(                  rMt*        j,                  j/                  | j(                  «      r$t        j                  d| j(                  › �«       �n| j(                  sX| j0                  › d| j                  › d�}t3        t3        | j4                  «      j6                  |«      j9                  «       | _        t        j                  d| j0                  › d| j(                  › d�«       t;        | «       nv| j(                  r<| j<                  r0t        j                  d| j(                  › d| j<                  › �«       n.t        j                  d| j0                  › d�«       t?        | «       d}| j@                  s‰| j                  t        j                  j                  k(  rb|r`|s|s|rZt        j                  d| j(                  › d�«       tC        | j(                  | jD                  «      }|st        jG                  d«       d}	d}
| jH                  s×|rÕ|s|s|rÏt        j                  d| j(                  › d�«       d | j                  › d�}t3        t3        | j4                  «      j6                  |«      j9                  «       }
tK        | j(                  |
| jD                  «      }	|	st        jG                  d!«       |	r+tM        | j(                  | jD                  «      st%        d"«      ‚|s| jN                  s|	rtt        j                  d#| j(                  › d�«       tQ        | j(                  | jD                  «       |	r/t        j                  d#|
› d�«       tQ        |
| jD                  «       |r,tS        jT                  | j0                  | jV                  ¬$«      }nf| j                   d%k(  r,tY        jT                  | j0                  | jV                  ¬$«      }n+t[        jT                  | j0                  | jV                  ¬$«      }| j\                  rt        j                  d&|› �«       |j^                  }|r|j^                  n|j`                  }|jb                  }| jb                  d'k7  r| jb                  }| j^                  d'k7  r| j^                  }| j`                  d'k7  r| j`                  }te        jf                  | j(                  d(¬)«      }| j                   › d*�|jh                  _5        d}| j                   dk(  rxtm        |jh                  | j                  «       |	rvte        jf                  |
d(¬)«      }| j                   › d+�|jh                  _5        tm        |jh                  | j                  «       n to        |jh                  | j                  «       d}|rg d,¢}n|s|rg d-¢}| jp                  r|js                  d.«       n|js                  d/«       | jt                  r|js                  d0«       n|js                  d/«       | jv                  r|js                  d1«       n|js                  d/«       |rX| jx                  r| jz                  r|js                  d2«       n|js                  d/«       | j|                  r|js                  d3«       d4g}| j                  r|js                  d5«       | j                   r$| j                  sJ d6«       ‚|js                  d7«       d}|r1td        j~                  j�                  d8||d9| j                   › �¬:«      }ne|r1td        j~                  j�                  d;||d<| j                   › �¬:«      }n2|r0td        j~                  j�                  d=||d>| j                   › �¬:«      }d?|_A        d}|rÈtd        j~                  j…                  d@|«      td        j~                  j…                  dA|«      td        j~                  j…                  dB| j†                  «      td        j~                  j…                  dC| jˆ                  rdnd«      td        j~                  j…                  dD| j                   dk(  rdnd«      g}�n/|r›td        j~                  j…                  d@|«      td        j~                  j…                  dA|«      td        j~                  j…                  dD| j                   dk(  rdnd«      td        j~                  j…                  dB| j†                  «      g}�n’|�r�td        j~                  j…                  d@|«      td        j~                  j…                  dA|«      td        j~                  j…                  dD| j                   dk(  rdnd«      td        j~                  j…                  dB| j†                  «      td        j~                  j…                  dE| jŠ                  «      td        j~                  j…                  dF| jŒ                  «      td        j~                  j…                  dG| jŽ                  «      td        j~                  j…                  dH| j�                  «      td        j~                  j…                  dI| jx                  «      td        j~                  j…                  dJ| j’                  «      g
}|r0|j•                  td        j~                  j…                  dK|«      g«       |j–                  j•                  |«       g }| j                   dLv �r”| jN                  rCt        j                  dM| j<                  › d�«       tQ        | j<                  | jD                  «       te        jf                  | j<                  d(¬)«      }t        |jh                  j˜                  «      dNk(  rdOndP}| j                   › dQ|› �|jh                  _5        t›        |jh                  | j                  «       t�        ||«       t�        ||«       |r¢| j"                  st%        dR«      ‚t        j                  dS«       tŸ        |jh                  «      rt        j                  dT«       nt        j                  dU«       t¡        |«      rt        j                  dV«       nt        j                  dW«       | j¢                  sHt¥        ||«      }t        j                  t        |«      › dX|D �cg c]  }|jj                  ‘Œ c}› dY�«       |j¦                  dk\  sJ dZ«       ‚|j–                  j•                  td        j~                  j…                  dO|jh                  «      td        j~                  j…                  d[|jh                  «      td        j~                  j…                  d\|j¦                  «      g«       �nÇ|	rò| j¢                  sHt¥        ||«      }t        j                  t        |«      › dX|D �cg c]  }|jj                  ‘Œ c}› d]�«       |r*t        j                  d^«       t©        |jh                  «       | j"                  r"t«        |jh                  |d«      st%        d_«      ‚|j–                  js                  td        j~                  j…                  d`|jh                  «      «       n6t­        |jh                  «      }t        j                  t        |«      › da�«       |r*t        j                  db«       t©        |jh                  «       | j"                  r"t«        |jh                  |d(«      st%        dc«      ‚|j–                  js                  td        j~                  j…                  d[|jh                  «      «       td        j~                  j¯                  ddt°        j²                  dedfg«      }td        j~                  j¯                  dgt°        j²                  dg«      }td        j~                  j¯                  dht°        j²                  dg«      }td        j~                  j¯                  dit°        j²                  dg«      }td        j~                  j¯                  djt°        j²                  dg«      }td        j~                  j¯                  dkt°        j´                  dg«      }td        j~                  j¯                  dlt°        j´                  dg«      } d}!|r
||||||| g}!n
|s|r|||| g}!| jp                  rAtd        j~                  j¯                  d.t°        j²                  |g«      }"|!js                  |"«       | jt                  rBtd        j~                  j¯                  d0t°        j²                  de|g«      }#|!js                  |#«       | jv                  rBtd        j~                  j¯                  d1t°        j²                  dedfg«      }$|!js                  |$«       | jx                  rN| jz                  rBtd        j~                  j¯                  d2t°        j²                  de|g«      }%|!js                  |%«       |rM| j|                  rAtd        j~                  j¯                  d3t°        j²                  dg«      }&|!js                  |&«       d}'|r2td        j~                  j¯                  d4t°        j²                  g dm¢«      }'n5|s|r1td        j~                  j¯                  d4t°        j²                  dedgg«      }'|'g}(| j                  rBtd        j~                  j¯                  d5t°        j´                  dedjg«      })|(js                  |)«       | j                   rDtd        j~                  j¯                  d7t°        j´                  dndedi|g«      }*|(js                  |*«       td        j~                  j·                  |g|s| j                   › do�n| j                   › dp�|!|(|«      }+td        j~                  j¹                  |+dq|jº                  ¬r«      },| jD                  rpddsl^m_}- |-jÁ                  td        jÂ                  «      |-jÁ                  dt«      k  rt        jG                  du«       tÅ        jÆ                  |,| j4                  d(d(¬v«       n te        jÆ                  |,| j4                  «       t        j                  dw| j4                  › �«       yc c}w c c}w )xzˆConvert model according to command line arguments.

    Args:
        args (argparse.Namespace): arguments parsed from command line
    r:   z**** past_present_share_buffer=r\   r   rI   )rÑ  rÏ  r�  r�  z**** Setting op_block_list to zI**** use --op_block_list if you want to override the block operator list.z<Currently only gpt2 with greedy search/sampling is supportedzLoutput_sequences_scores currently is not supported in greedy search/samplingzHoutput_token_scores currently is not supported in greedy search/samplingzi`use_decoder_masked_attention` MUST be turned on to use `past_present_share_buffer` in case of BeamSearchzS`past_present_share_buffer` MUST be turned on to use `use_decoder_masked_attention`z?`use_decoder_masked_attention` option is only supported on GPUsz)skip convert_to_onnx since path existed: Ú_past_z.onnxzConvert GPT model z	 to onnx z ...z,skip convert_to_onnx since paths specified: z and zConvert model z to onnx ...Fz=Pad logits MatMul weights for optimal MatMul perf in fp16 on z. The file will be overwritten.z]Tried and failed to pad logits MatMul weights. Performance may be sub-optimal for this MatMulNz*Creating an initial run GPT2 decoder from z. Úgpt2_init_past_zuTried and failed to generate the init decoder GPT2 model. Performance may be sub-optimal for the initial decoding runzGCould not update the input shapes for the non-initial decoder subgraph.z Run symbolic shape inference on ©r†   r;   zConfig=rb   Tr¤   z decoderz init decoder©rí   Ú
max_lengthÚ
min_lengthÚ	num_beamsÚnum_return_sequencesÚlength_penaltyÚrepetition_penalty©rí   r½  r¾  rÂ  rU   r   rX   rï   rZ   r[   Ú	sequencesÚsequences_scoresz8--output_token_scores requires --output_sequences_scoresÚscoresÚ
BeamSearchÚBeamSearch_râ  ÚGreedySearchÚGreedySearch_ÚSamplingÚ	Sampling_rÿ  Úeos_token_idÚpad_token_idÚno_repeat_ngram_sizerT   r˜   ÚtemperatureÚtop_pÚfilter_valueÚmin_tokens_to_keepÚcustomÚpresence_penaltyÚ
vocab_size©r;   r<   zSymbolic shape inference on r¶   rD  zencoder and decoder initú zMpast_present_share_buffer is only supported with use_decoder_masked_attentionzl*****update t5 decoder subgraph to share past/present buffer and use decoder_masked_multihead_attention*****z4*****update t5 decoder subgraph successfully!!!*****zF*****DecoderMaskedMultiHeadAttention is not applied to T5 decoder*****z9*****pack qkv for decoder masked mha successfully!!!*****z3*****pack qkv for decoder masked mha failed!!!*****z shared initializers (z>) in encoder and decoder subgraphs are moved to the main graphz%decoder_start_token_id should be >= 0rE  Údecoder_start_token_idzC) in decoder and init decoder subgraphs are moved to the main graphzY*****update init decoder subgraph to make past and present share buffer******************zLCould not update the init decoder subgraph to use DecoderMaskedSelfAttentionÚinit_decoderz: initializers from the decoder are moved to the main graphzT*****update decoder subgraph to make past and present share buffer******************zGCould not update the decoder subgraph to use DecoderMaskedSelfAttentionrí   rƒ  rÃ  r½  r¾  r¿  rÀ  rÁ  rÂ  )rƒ  rÀ  r½  zmax_length - sequence_lengthz beam searchz greedy searchzonnxruntime.transformers)Úproducer_nameÚopset_imports)Úversionz1.12.0z0Require onnx >= 1.12 to save large (>2GB) model!)r©   Úall_tensors_to_one_filezmodel save to )dr˜   r    r.   r/   r0   rV   r‹   rŒ   r‰   rŠ   r…   r   rt   r&   ÚNotImplementedErrorrR   rS   rW   rû   re   r„   rp   rq   Úexistsrƒ   r   r›   rœ   Úas_posixr�   rž   r    rN   rÛ   rL   r¯   rO   r®  rŒ  rM   r²   r   Úfrom_pretrainedr†   r   r   rD   rÍ  rÎ  rÖ  r«   r¬   r¼   r½   r  r  rU   rˆ   rX   rY   rÔ  rZ   r[   r0  rt  r^  rO  rÏ  rT   rÐ  rÑ  rÒ  rÓ  rÕ  r‡   r[  rÁ   r   r¶  rp  r†  rP   rG  rÙ  r|  r�  rK  r1  r   rý   rþ   Ú
make_graphÚ
make_modelÚopset_importÚ	packagingrÝ  ÚparseÚ__version__r   r®   ).r~   r·  Úis_gpt2Úis_beamsearchÚis_greedysearchÚis_samplingrV   Úonnx_filenameÚlogits_matmul_weight_paddedÚgpt2_init_decoder_generatedÚgpt2_init_decoder_onnx_pathÚgpt2_init_decoder_onnx_filenamer²  rÍ  rÎ  rÖ  rÑ   rš  rã  rä  r+  Úattr_to_extendrF  r=  Úsuffixr  rí   r½  r¾  r¿  rÀ  rÁ  rÂ  Úgraph_inputsrU   rX   rï   rZ   r[   rÄ  Úgraph_outputsrÅ  rÆ  Ú	new_graphÚ	new_modelrÝ  s.                                                 r(   Úconvert_generation_modelrø  û	  sé  € ð —O‘O vÑ-€GØ)¬^×-FÑ-FÑF€MØ+¬~×/JÑ/JÑJ€OØ'¬>×+BÑ+BÑB€KØ&*×&DÑ&DÐä
‡K�KÐ1Ð2KÐ1LÐMÔNÜ
ˆ4×ÑÓ !Ò#¨×(:Ñ(:¸1Ñ(=ÀÒ(GÙ�t—~‘~¬×):Ñ):×)@Ñ)@Ò@ò"ˆDÔô �K‰KÐ8¸×9KÑ9KÐ8LÐMÔNÜ�K‰KÐcÕdà!#ˆDÔá™+ÙÜ%Ð&dÓeÐeØ×'Ò'Ü%Ð&tÓuÐuØ×#Ò#Ü%Ð&pÓqÐqñ !¡]¸4×;\Ò;\ÜØwó
ð 	
ð ×(Ò(Ñ1JÜÐnÓoÐoð ×(Ò(°·²ÜÐZÓ[Ð[áØ×Ò¤§¡§¡°×0AÑ0AÔ!BÜ�K‰KÐCÀD×DUÑDUÐCVÐWÖXà×$Ò$Ø#'×#:Ñ#:Ð";¸6À$Ç.Á.ÐAQÐQVÐ W�Ü$(¬¨d¯k©kÓ):×)AÑ)AÀ=Ó$Q×$ZÑ$ZÓ$\�Ô!ä�K‰KÐ,¨T×-DÑ-DÐ,EÀYÈt×O`ÑO`ÐNaÐaeÐfÔgÜ˜Õà×Ò ×!?Ò!?Ü�K‰KØ>¸t×?PÑ?PÐ>QÐQVÐW[×WuÑWuÐVvÐwõô �K‰K˜.¨×)@Ñ)@Ð(AÀÐNÔOÜ�tÔð #(Ðà×'Ò'Ø�N‰Nœi×/Ñ/×5Ñ5Ò5ÙÙ™o±ä�‰ØKÈD×L]ÑL]ÐK^ð _,ð ,ô	
ô 'CÀ4×CTÑCTÐVZ×VsÑVsÓ&tÐ#Ù*Ü�N‰NØoôð #(ÐØ"&Ðà×;Ò;ÙÙ™o±ä�‰Ð@À×ARÑARÐ@SÐSUÐVÔWà,;¸D¿N¹NÐ;KÈ5Ð*QÐ'ä&*¬4°·±Ó+<×+CÑ+CÐEdÓ&e×&nÑ&nÓ&pÐ#ä&@Ø×ÑØ'Ø×)Ñ)ó'
Ð#ñ +Ü�N‰NðNôñ 'Ô/YØ×Ñ˜t×<Ñ<ô0
ô ÐfÓgÐgñ
 # d×&>Ò&>ÑB]Ü�‰Ð6°t×7HÑ7HÐ6IÐIhÐiÔjÜ˜×)Ñ)¨4×+HÑ+HÔIÙ&Ü�K‰KÐ:Ð;VÐ:WÐWvÐwÔxÜÐ7¸×9VÑ9VÔWáÜ×+Ñ+¨D×,CÑ,CÈtÏ~É~Ô^‰Ø	�‰˜DÒ	 Ü×)Ñ)¨$×*AÑ*AÈTÏ^É^Ô\‰ä×*Ñ*¨4×+BÑ+BÈdÏnÉnÔ]ˆà‡|‚|Ü�‰�g˜f˜XÐ&Ô'à×&Ñ&€LÙ*1�6×&Ò&°v×7JÑ7J€LØ×"Ñ"€Jð ‡�˜"ÒØ—_‘_ˆ
à×Ñ˜BÒØ×(Ñ(ˆØ×Ñ˜BÒØ×(Ñ(ˆä—O‘O D×$5Ñ$5È$ÔO€MØ"&§/¡/Ð!2°(Ð;€M×ÑÔà"ÐØ‡�˜&Ò Ü˜]×0Ñ0°$·.±.ÔAñ 'Ü&*§o¡oÐ6QÐfjÔ&kÐ#Ø48·O±OÐ3DÀMÐ1RÐ#×)Ñ)Ô.Ü Ð!8×!>Ñ!>ÀÇÁÕOä" =×#6Ñ#6¸¿¹ÔGà€FÙò
‰ñ 
™Kò
ˆð ‡‚Ø�‰�lÕ#à�‰�bÔà×ÒØ�‰Ð)Õ*à�‰�bÔà×!Ò!Ø�‰Ð&Õ'à�‰�bÔáØ�;Š;˜4×-Ò-Ø�M‰M˜/Õ*à�M‰M˜"Ôà�9Š9Ø�M‰M˜&Ô!àˆm€GØ×#Ò#Ø�‰Ð)Ô*à×ÒØ×+Ò+ÐgÐ-gÓgÐ+Ø�‰�xÔ à€DÙÜ�{‰{×$Ñ$ØØØØ˜tŸ™Ð/Ð0ð	 %ó 
‰ñ 
Ü�{‰{×$Ñ$ØØØØ  §¡Ð 1Ð2ð	 %ó 
‰ñ 
Ü�{‰{×$Ñ$ØØØØ˜TŸ_™_Ð-Ð.ð	 %ó 
ˆð "€D„Kà€NÙä�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ&Ð'=¸t×?XÑ?XÓYÜ�K‰K×&Ñ&Ð'7¸d×>QÒ>Q¹ÐWXÓYÜ�K‰K×&Ñ& |¸$¿/¹/ÈVÒ:S±QÐYZÓ[ð
Šñ 
ä�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ& |¸$¿/¹/ÈVÒ:S±QÐYZÓ[Ü�K‰K×&Ñ&Ð'=¸t×?XÑ?XÓYð	
Šò 
ä�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ& ~°|ÓDÜ�K‰K×&Ñ& |¸$¿/¹/ÈVÒ:S±QÐYZÓ[Ü�K‰K×&Ñ&Ð'=¸t×?XÑ?XÓYÜ�K‰K×&Ñ& }°d×6FÑ6FÓGÜ�K‰K×&Ñ& w°·
±
Ó;Ü�K‰K×&Ñ& ~°t×7HÑ7HÓIÜ�K‰K×&Ñ&Ð';¸T×=TÑ=TÓUÜ�K‰K×&Ñ& x°·±Ó=Ü�K‰K×&Ñ&Ð'9¸4×;PÑ;PÓQð
ˆñ #Ø×ÑœtŸ{™{×9Ñ9¸,È
ÓSÐTÔUà‡N�N×Ñ˜.Ô)à€Là‡�˜-Ò'Ø×#Ò#Ü�K‰KÐ6°t×7UÑ7UÐ6VÐVuÐvÔwÜ˜D×:Ñ:¸D×<YÑ<YÔZÜŸ™¨×(FÑ(FÐ[_Ô`ˆÜ! -×"5Ñ"5×";Ñ";Ó<ÀÒA‘ÐGaˆØ&*§o¡oÐ%6°a¸°xÐ#@ˆ×ÑÔ Ü/°×0CÑ0CÀTÇ^Á^ÔTä! -°Ô8Ü! -°Ô8ñ %Ø×4Ò4Ü Ð!pÓqÐqä�K‰KØ~ôô OÈ}×ObÑObÔcÜ—‘ÐRÕSä—‘ÐdÔeä.¨}Ô=Ü—‘ÐWÕXä—‘ÐQÔRà×/Ò/ä2°=À-ÓPˆLÜ�K‰KÜ�|Ó$Ð%Ð%;È\Ö<ZÈ¸Q¿V»VÒ<ZÐ;[ð  \Zð  [ôð ×,Ñ,°Ò1ÐZÐ3ZÓZÐ1à�‰×Ñä—‘×*Ñ*¨9°m×6IÑ6IÓJÜ—‘×*Ñ*¨9°m×6IÑ6IÓJÜ—‘×*Ñ*Ð+CÀV×EbÑEbÓcðö	
ñ 'ð ×3Ò3ä6Ð7NÐP]Ó^�Ü—‘Ü˜<Ó(Ð)Ð)?ÐQ]Ö@^ÈAÀÇÃÒ@^Ð?_ð  `cð  dôñ
 )Ü—‘ÐwÔxÜAÐBY×B_ÑB_Ô`ð ×0Ò0Ô9mØ'×-Ñ-¨}¸eô:ô !Ð!oÓpÐpà�N‰N×!Ñ!¤$§+¡+×"<Ñ"<¸^ÐMd×MjÑMjÓ"kÕlô -¨]×-@Ñ-@ÓAˆLÜ�K‰Kœ3˜|Ó,Ð-Ð-gÐhÔiñ %Ü�K‰KÐnÔoÜ=¸m×>QÑ>QÔRð ×,Ò,Ô5iØ×Ñ °ô6
ô ÐfÓgÐgà�‰×ÑœdŸk™k×8Ñ8¸ÀM×DWÑDWÓXÔYô —‘×2Ñ2°;Ä×@QÑ@QÐT`ÐbsÐStÓu€IÜ—‘×3Ñ3°LÄ+×BSÑBSÐVWÐUXÓY€JÜ—‘×3Ñ3°LÄ+×BSÑBSÐVWÐUXÓY€JÜ—‘×2Ñ2°;Ä×@QÑ@QÐTUÐSVÓW€IÜŸ;™;×=Ñ=Ð>TÔVa×VgÑVgÐjkÐilÓmÐÜ—[‘[×7Ñ7Ð8HÌ+×J[ÑJ[Ð^_Ð]`Óa€NÜŸ™×;Ñ;Ð<PÔR]×RcÑRcÐfgÐehÓiÐà€LÙàØØØØ ØØð
‰ñ 
™KàØØØð	
ˆð ‡‚Ü—[‘[×7Ñ7¸Äk×FWÑFWÐZdÐYeÓfˆ
Ø×Ñ˜JÔ'à×ÒÜ ŸK™K×>Ñ>Ø¤×!2Ñ!2°\À:Ð4Nó
Ðð 	×ÑÐ-Ô.à×!Ò!ÜŸ™×;Ñ;Øœk×/Ñ/°,Ð@QÐ1Ró
ˆð 	×Ñ˜NÔ+à‡{‚{�t×)Ò)ÜŸ™×:Ñ:Øœ[×.Ñ.°¸zÐ0Jó
ˆð 	×Ñ˜MÔ*á�t—y’yÜ�{‰{×1Ñ1°&¼+×:KÑ:KÈaÈSÓQˆØ×Ñ˜DÔ!ð €IÙÜ—K‘K×6Ñ6ØÜ×ÑÚ@ó
‰	ñ
 
™KÜ—K‘K×6Ñ6ØÜ×ÑØ˜<Ð(ó
ˆ	ð �K€Mà×#Ò#ÜŸ;™;×=Ñ=ØÜ×ÑØÐ1Ð2ó
Ðð
 	×ÑÐ-Ô.à×ÒÜ—‘×3Ñ3ØÜ×ÑØ+¨\¸;È
ÐSó
ˆð
 	×Ñ˜VÔ$ä—‘×&Ñ&Ø	ˆÙ1@ˆD�O‰OÐ˜LÑ	)ÈÏÉÐHYÐYgÐFhØØØó€Iô —‘×&Ñ&ØØ0Ø#×0Ñ0ð 'ó €Ið ×$Ò$Ý%à�=‰=œ×)Ñ)Ó*¨W¯]©]¸8Ó-DÒDÜ�N‰NÐMÔNä�‰ØØ�K‰KØ"&Ø$(ö		
ô 	�	‰	�)˜TŸ[™[Ô)Ü
‡K�K�. §¡ Ð.Õ/ùòk =[ùò8 A_s   ø?AXýAX
rí   rï   rÍ  rÎ  Úbad_words_idsc                 ó°  — | j                   r)t        j                  j                  «       st	        d«      ‚| j
                  t        j                  j                  k(  r|j                  «        t        j                  | j                   rdnd«      }|j                  |«       t        j                  d«       |j                  |«      }|j                  |«      }g }t        | j                  «      D ]È  }	t        j                  «       }
|j!                  ||| j"                  | j$                  | j&                  | j(                  | j*                  ||| j,                  | j.                  | j0                  |r|ndd| j2                  xs | j4                  ¬«      }	|j7                  t        j                  «       |
z
  «       ŒÊ |j8                  d   }dd	lm}  |||«      S )
aœ  Test PyTorch performance of text generation.

    Args:
        args (argparse.Namespace): arguments parsed from command line
        model (Union[GPT2LMHeadModel, T5ForConditionalGeneration]): PyTorch model
        input_ids (torch.Tensor): input_ids
        attention_mask (torch.Tensor): Attention mask
        eos_token_id (int): EOS token ID
        pad_token_id (int): Padding token ID
        bad_words_ids (List[List[int]]): Words shall not be generated.

    Raises:
        RuntimeError: PyTorch with CUDA is not available for --use_gpu

    Returns:
        Dict[str, Any]: A dictionary with string with metric name, and value can be integer or string.
    z=Please install PyTorch with Cuda for testing gpu performance.zcuda:0ÚcpuFNT©rí   rï   r½  r¾  r¿  rT   rÏ  rÍ  rÎ  rÀ  rÁ  rÂ  rù  Úreturn_dict_in_generateÚoutput_scoresr   ©Úget_latency_result)re   ÚtorchÚcudaÚis_availablerä   r…   r   rt   r&   ÚhalfÚdevicerç  Úset_grad_enabledrú   Ú
total_runsÚtimeÚgenerater½  r¾  r¿  rT   rÏ  rÀ  rÁ  rÂ  rR   rS   rˆ   r/  Úbenchmark_helperr   )r~   r°   rí   rï   rÍ  rÎ  rù  r  Útorch_latencyrÂ  Ústartrƒ  r   s                r(   Útest_torch_performancer  ÷  sw  € ð4 ‡|‚|œEŸJ™J×3Ñ3Ô5ÜÐZÓ[Ð[à‡~�~œ×*Ñ*×0Ñ0Ò0Ø�
‰
Œä�\‰\ d§l¢l™(¸Ó>€FØ	‡H�HˆVÔä	×Ñ˜5Ô!Ø—‘˜VÓ$€IØ#×&Ñ& vÓ.€Nà€MÜ�4—?‘?Ó#ò 2ˆÜ—	‘	“ˆØ�N‰NØØ)Ø—‘Ø—‘Ø—n‘nØ×.Ñ.Ø!%×!:Ñ!:Ø%Ø%Ø!%×!:Ñ!:Ø×.Ñ.Ø#×6Ñ6Ù+8™-¸dØ$(Ø×6Ñ6ÒR¸$×:RÑ:Rð ó 
ˆð" 	×ÑœTŸY™Y›[¨5Ñ0Õ1ð'2ð( —‘ Ñ#€JÝ3á˜m¨ZÓ8Ð8r*   c                 ó  — t        j                  | j                  t         j                  ¬«      }t	        | j                  d   «      D ]?  }d}t	        | j                  d   «      D ]   }| |   |   |k(  r|dk(  r	d||   |<   Œ|dz  }Œ" ŒA |S )Nr¸   r   r\   )rÉ   Úonesr/  Úint32rú   )rí   rÎ  rï   r  Úabs_posr:  s         r(   Úcreate_attention_maskr  9  s�   € Ü—W‘W˜YŸ_™_´B·H±HÔ=€NÜ�9—?‘? 1Ñ%Ó&ò ˆØˆÜ�y—‘ qÑ)Ó*ò 	ˆAØ˜‰|˜A‰ ,Ò.°7¸a²<Ø'(�˜qÑ! !Ò$à˜1‘‘ñ		ðð Ðr*   Ú	sentencesÚ	is_greedyc                 óæ  — | j                   dk(  sJ ‚t        j                  | j                  | j                  ¬«      }d|_        |j                  |_        t        j                  | j                  | j                  |j                  ¬«      }|€g d¢} ||dd¬	«      }|d
   }|d   }d}|j                  |d¬«      }	|	D �
cg c]  }
|
g‘Œ }	}
| j                  rt        j                  d|	«       ng }	|j                  }|j                  }|j                  }|j                  }g }d}| j                   �sdt#        d«       t#        d«       |j%                  ||| j&                  | j(                  | j*                  | j,                  | j.                  ||| j0                  | j2                  | j4                  |	r|	ndd| j6                  xs | j8                  ¬«      }t#        d
|«       t#        d«       t#        d|j:                  «       | j6                  rt#        d|j<                  «       | j8                  rt#        d|j>                  «       tA        |j:                  «      D ]9  \  }}|jC                  |d¬«      }|jE                  |«       t#        |› d|› �«       Œ; t#        d«       t#        d«       |rÌ|jG                  «       jI                  «       jK                  tL        jN                  «      tM        jP                  | j&                  gtL        jN                  ¬«      tM        jP                  | j(                  gtL        jN                  ¬«      tM        jP                  | j4                  gtL        jR                  ¬«      dœ}�nW|jG                  «       jI                  «       jK                  tL        jN                  «      tM        jP                  | j&                  gtL        jN                  ¬«      tM        jP                  | j(                  gtL        jN                  ¬«      tM        jP                  | j*                  gtL        jN                  ¬«      tM        jP                  | j0                  gtL        jN                  ¬«      tM        jP                  | j2                  gtL        jR                  ¬«      tM        jP                  | j4                  gtL        jR                  ¬«      dœ}| j                  rBtM        jT                  |tL        jN                  ¬«      }| j                  r|	D ]  }d||<   Œ	 ||d<   | jV                  rtY        ||«      |d<   |jZ                  d   }| j\                  rAt        j_                  d«       tM        jT                  ||ftL        jN                  ¬«      }||d<   | j`                  r­tc        | jd                  «      jf                  ji                  «       }t        j                  d |«       dd!l5m6} t        j_                  d"|› d#�«       |g}tA        |«      D ]:  \  }}tn        jp                  js                  |d$tu        |«      z   «      } |||«       Œ< t        j                  d%|«       | jv                  ryt        j                  d&«       ty        | jd                  | jz                  | j|                  «      }t        j                  d'«       |j                  d|«      }g }t�        | j‚                  «      D ]N  }t…        j„                  «       } |j                  d|«      }|jE                  t…        j„                  «       | z
  «       ŒP dd(lCmD}! |jZ                  d   } |!||«      }"t#        d)«       |d   }#t#        d|#«       | j6                  rt#        d|d*   «       | j8                  rt#        d|d+   «       |rZ|#jZ                  \  }}$g }%t�        |«      D ]:  }|jC                  |#|   d¬«      }|%jE                  |«       t#        d,|› d-|› �«       Œ< np|#jZ                  \  }}&}$g }%t�        |«      D ]P  }t�        |&«      D ]@  }'|jC                  |#|   |'   d¬«      }|%jE                  |«       t#        d,|› d.|'› d|› �«       ŒB ŒR |r¹|j:                  j‹                  || j0                  d/«      }(t�        jŽ                  |#«      })t#        d«       t#        d0«       t#        |(«       t#        |«       t#        d«       t#        d1«       t#        |)«       t#        |%«       t#        d«       ||%k(  }*t#        d2|*rd3nd4«       |*|"d5<   | j�                  rt“        | ||||||	«      }+t#        d6|+«       t#        d7|"«       |"S c c}
w )8a9  Test GPT-2 model

    Args:
        args (argparse.Namespace): arguments parsed from command line
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    r:   r»  Úleft)r†   rÎ  N)zThe product is releasedzI enjoy walking in the parkzTest best way to investÚptT©Úreturn_tensorsrØ   rí   rï   úwalk in park)Úadd_prefix_spacerù  ú2--------------------------------------------------úCTest PyTorch model and beam search with huggingface transformers...rü  ú!huggingface transformers outputs:rÄ  rÅ  rÆ  ©Úskip_special_tokensú: ú'Testing beam search with onnxruntime...r¸   rÃ  r¼  r   rU   zYUse prefix vocab mask with all ones in ORT, but no corresponding setting for Torch model.rX   Útest_data_dir©Úoutput_test_datazSaving test_data to z/test_data_set_* ...Útest_data_set_ú
ORT inputszCreating ort session......zRun ort session......rÿ  úORT outputs:r\   r¶   úbatch z sequence: ú
 sequence rb   úTorch Sequences:úORT Sequences:zTorch and ORT result isÚsameÚ	differentÚparityúTorch LatencyÚORT)Jr˜   r   râ  rƒ   r†   Úpadding_sideÚ	eos_tokenÚ	pad_tokenr   rÍ  ÚencoderU   r‹   r�   r²  rÖ  rf   rN  r	  r½  r¾  r¿  rT   rÏ  rÀ  rÁ  rÂ  rR   rS   rÄ  rÅ  rÆ  rü   Údecoderˆ   rû  ÚnumpyÚastyperÉ   r  ÚarrayÚfloat32r  rY   r  r/  rX   rŒ   ri   r   r›   rœ   rá  Úbert_test_datar%  rp   rq   ro   rn   rg   rê   re   rc   Úrunrú   r  r  r
  r   r%  r  Ú
LongTensorrh   r  ),r~   r  r  Ú	tokenizerr°   rã  rí   rï   Ú	bad_wordsrù  Úword_idr²  rÍ  rÎ  rÖ  Útorch_decoded_sequencesÚbeam_outputsr  ÚsequenceÚdecoded_sequencerU   Úbad_word_idrƒ  rX   r#  r%  Ú
all_inputsÚdirré   ÚresultÚlatencyrÂ  r  r   r›   rÄ  r½  Úort_decoded_sequencesÚnum_sequencesr:  Útorch_sequencesÚort_sequencesÚis_sameÚtorch_latency_outputs,                                               r(   Útest_gpt_modelrP  E  s  € ð �?‰?˜fÒ$Ð$Ð$ä×-Ñ-¨d×.EÑ.EÐQU×Q_ÑQ_Ô`€IØ#€IÔØ#×-Ñ-€IÔä×+Ñ+Ø×ÑØ—.‘.Ø×+Ñ+ô€Eð Ðò
ˆ	ñ �y°¸tÔD€FØ�{Ñ#€IØÐ,Ñ-€Nà€IØ×$Ñ$ YÀÐ$ÓF€MØ.;Ö< 7�g’YÐ<€MÐ<Ø‡‚Ü�‰�_ mÕ4àˆà�\‰\€FØ×&Ñ&€LØ×&Ñ&€LØ×"Ñ"€Jà ÐØ€LØ×ÓÜˆhŒÜÐSÔTØ—~‘~ØØ)Ø—‘Ø—‘Ø—n‘nØ×.Ñ.Ø!%×!:Ñ!:Ø%Ø%Ø!%×!:Ñ!:Ø×.Ñ.Ø#×6Ñ6Ù+8™-¸dØ$(Ø×6Ñ6ÒR¸$×:RÑ:Rð &ó 
ˆô" 	ˆk˜9Ô%ÜÐ1Ô2Üˆk˜<×1Ñ1Ô2Ø×'Ò'ÜÐ$ l×&CÑ&CÔDØ×#Ò#Ü�(˜L×/Ñ/Ô0Ü$ \×%;Ñ%;Ó<ò 	.‰KˆAˆxØ(×/Ñ/°ÈdÐ/ÓSÐØ#×*Ñ*Ð+;Ô<Ü�Q�C�rÐ*Ð+Ð,Õ-ð	.ô
 
ˆ(„OÜ	Ð
3Ô4áà"Ÿ™›×.Ñ.Ó0×7Ñ7¼¿¹ÓAÜŸ(™( D§O¡OÐ#4¼B¿H¹HÔEÜŸ(™( D§O¡OÐ#4¼B¿H¹HÔEÜ"$§(¡(¨D×,CÑ,CÐ+DÌBÏJÉJÔ"Wñ	
Šð #Ÿ™›×.Ñ.Ó0×7Ñ7¼¿¹ÓAÜŸ(™( D§O¡OÐ#4¼B¿H¹HÔEÜŸ(™( D§O¡OÐ#4¼B¿H¹HÔEÜŸ™ 4§>¡>Ð"2¼"¿(¹(ÔCÜ$&§H¡H¨d×.GÑ.GÐ-HÔPR×PXÑPXÔ$YÜ Ÿh™h¨×(;Ñ(;Ð'<ÄBÇJÁJÔOÜ"$§(¡(¨D×,CÑ,CÐ+DÌBÏJÉJÔ"Wñ
ˆð ‡‚Ü—W‘W˜j´·±Ô:ˆ
Ø�?Š?Ø,ò ,�Ø*+�
˜;Ò'ð,à)ˆˆ|Ñà×!Ò!Ü#8¸ÀLÓ#QˆÐÑ à—‘ Ñ#€JØ×ÒÜ�‰ÐoÔpÜŸG™G Z°Ð$<ÄBÇHÁHÔMÐØ&7ˆÐ"Ñ#à×ÒÜ˜TŸ[™[Ó)×0Ñ0×9Ñ9Ó;ˆÜ�‰�_ mÔ4Ý3ä�‰Ð*¨=¨/Ð9MÐNÔOà�Xˆ
Ü" :Ó.ò 	*‰IˆAˆvÜ—'‘'—,‘,˜}Ð.>ÄÀQÃÑ.GÓHˆCÙ˜S &Õ)ð	*ô ‡L�L�˜vÔ&à×ÒØä
‡L�LÐ-Ô.Ü$ T§[¡[°$·,±,À×@XÑ@XÓY€Kä
‡L�LÐ(Ô)Ø�_‰_˜T 6Ó*€Fð €GÜ�4—?‘?Ó#ò ,ˆÜ—	‘	“ˆØ�O‰O˜D &Ó)ˆØ�‰”t—y‘y“{ UÑ*Õ+ð,õ
 4à—‘ Ñ#€JÙ ¨Ó4€Fä	ˆ.ÔØ�q‘	€IÜ	ˆ+�yÔ!Ø×#Ò#ÜÐ  &¨¡)Ô,Ø×ÒÜˆh˜˜q™	Ô"áØ#,§?¡?Ñ ˆ�ZØ "ÐÜ�zÓ"ò 	=ˆAØ(×/Ñ/°	¸!±ÐRVÐ/ÓWÐØ!×(Ñ(Ð)9Ô:Ü�F˜1˜#˜[Ð)9Ð(:Ð;Õ<ñ	=ð
 3<·/±/Ñ/ˆ�] JØ "ÐÜ�zÓ"ò 	EˆAÜ˜=Ó)ò E�Ø#,×#3Ñ#3°I¸a±LÀ±OÐY]Ð#3Ó#^Ð Ø%×,Ñ,Ð-=Ô>Ü˜˜q˜c ¨A¨3¨bÐ1AÐ0BÐCÕDñEð	Eñ Ø&×0Ñ0×8Ñ8¸ÀT×E^ÑE^Ð`bÓcˆÜ×(Ñ(¨Ó3ˆÜˆhŒÜÐ Ô!ÜˆoÔÜÐ%Ô&ÜˆhŒÜÐÔÜˆmÔÜÐ#Ô$ÜˆhŒà)Ð-BÑBˆÜÐ'±7©ÀÔLØ"ˆˆxÑà×ÒÜ5ØØØØØØØó 
Ðô 	ˆoÐ3Ô4ä	ˆ%�Ôà€MùòY =s   Â?
e.c                 óä  — | j                   dv sJ ‚| j                  rt        j                  d«       yt	        j
                  | j                  | j                  ¬«      }d|_        | j                   dk(  r,t        j
                  | j                  | j                  ¬«      }n+t        j
                  | j                  | j                  ¬«      }|€ddg} ||d	d
¬«      }|d   }|d   }d}|j                  |«      dd }|D �	cg c]  }	|	g‘Œ }}	| j                  rt        j                  d|«       ng }|j                  }
|
j                  }|
j                  }|
j                   }t        j                  d|› d|› d|› �«       g }| j"                  �sdt%        d«       t%        d«       |j'                  ||| j(                  | j*                  | j,                  | j.                  | j0                  ||| j2                  | j4                  | j6                  |r|ndd
| j8                  xs | j:                  ¬«      }t%        d|«       t%        d«       t%        d|j<                  «       | j8                  rt%        d|j>                  «       | j:                  rt%        d|j@                  «       tC        |j<                  «      D ]9  \  }}|jE                  |d
¬«      }|jG                  |«       t%        |› d|› �«       Œ; t%        d«       t%        d«       tI        jJ                  |tH        jL                  ¬«      }| j                  r|D ]  }d||<   Œ	 |jO                  «       jQ                  «       jS                  tH        jL                  «      tI        jT                  | j(                  gtH        jL                  ¬«      tI        jT                  | j*                  gtH        jL                  ¬«      tI        jT                  | j,                  gtH        jL                  ¬«      tI        jT                  | j2                  gtH        jL                  ¬«      tI        jT                  | j4                  gtH        jV                  ¬«      tI        jT                  | j6                  gtH        jV                  ¬«      d œ}| j                  r||d!<   | jX                  rt[        ||«      |d<   | j\                  r”t_        | j`                  «      jb                  je                  «       }t        j                  d"|«       dd#l3m4} |g}tC        |«      D ]:  \  }}tj        jl                  jo                  |d$tq        |«      z   «      } |||«       Œ< t        j                  d%|«       ts        | j`                  | jt                  | jv                  «      }g }ty        | jz                  «      D ]N  }t}        j|                  «       }|j                  d|«      }|jG                  t}        j|                  «       |z
  «       ŒP |j€                  d   }dd&lAmB}  |||«      } t%        d'«       d   }!t%        d|!«       | j8                  rt%        d|d(   «       | j:                  rt%        d|d)   «       |!j€                  \  }}"}#g }$ty        |«      D ]P  }ty        |"«      D ]@  }%|jE                  |!|   |%   d
¬«      }|$jG                  |«       t%        d*|› d+|%› d|› �«       ŒB ŒR | j"                  s¹j<                  j‡                  || j2                  d«      }&t‰        jŠ                  |!«      }'t%        d«       t%        d,«       t%        |&«       t%        |«       t%        d«       t%        d-«       t%        |'«       t%        |$«       t%        d«       ||$k(  }(t%        d.|(rd/nd0«       |(| d1<   | jŒ                  rt�        | ||||||«      })t%        d2|)«       t%        d3| «       | S c c}	w )4a=  Test T5 or MT5 model

    Args:
        args (argparse.Namespace): arguments parsed from command line
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    r×  zLSkipping parity test as prefix vocab mask is not implemented by Hugging FaceNr»  r  r;   z4translate English to French: The product is releasedz¬summarize: research continues to show that pets bring real health benefits to their owners. Having a dog around can lead to lower levels of stress for both adults and kids.r  Tr  rí   rï   r  rb   rù  zeos_token_id:z, pad_token_id:z, vocab_size:r  r  rü  r  rÄ  rÅ  rÆ  r  r!  r"  r¸   r   r¼  rU   r#  r$  r&  r'  rÿ  r(  r\   r¶   r)  r*  r+  r,  zTorch and ORT result is r-  r.  r/  r0  r1  )Hr˜   rX   r‹   r�   r   râ  rƒ   r†   r2  r   r   r5  rU   r²  rÍ  rÎ  rÖ  rf   rN  r	  r½  r¾  r¿  rT   rÏ  rÀ  rÁ  rÂ  rR   rS   rÄ  rÅ  rÆ  rü   r6  rˆ   rÉ   r  r  rû  r7  r8  r9  r:  rY   r  ri   r   r›   rœ   rá  r;  r%  rp   rq   ro   rn   rê   re   rc   rú   r  r  r<  r/  r
  r   r%  r  r=  rh   r  )*r~   r  r>  r°   rã  rí   rï   r?  rù  r@  r²  rÍ  rÎ  rÖ  rA  rB  r  rC  rD  rU   rE  r#  r%  rF  rG  ré   rI  rÂ  r  rH  rƒ  r   r›   rÄ  rK  r½  rJ  r:  rL  rM  rN  rO  s*                                             r(   Útest_t5_modelrR    sá  € ð �?‰?˜mÑ+Ð+Ð+à×ÒÜ�‰ÐcÔdØä×+Ñ+¨D×,CÑ,CÈtÏ~É~Ô^€IØ#€IÔà‡�˜$ÒÜ*×:Ñ:Ø×#Ñ#Ø—n‘nô
‰ô
 ,×;Ñ;Ø×#Ñ#Ø—n‘nô
ˆð ÐàBð {ð
ˆ	ñ �y°¸tÔD€FØ�{Ñ#€IØÐ,Ñ-€Nà€IØ×$Ñ$ YÓ/°°Ð4€MØ.;Ö< 7�g’YÐ<€MÐ<Ø‡‚Ü�‰�_ mÕ4àˆà�\‰\€FØ×&Ñ&€LØ×&Ñ&€LØ×"Ñ"€JÜ
‡L�L�=  ¨o¸l¸^È=ÐYcÐXdÐeÔfà ÐØ×ÓÜˆhŒÜÐSÔTØ—~‘~ØØ)Ø—‘Ø—‘Ø—n‘nØ×.Ñ.Ø!%×!:Ñ!:Ø%Ø%Ø!%×!:Ñ!:Ø×.Ñ.Ø#×6Ñ6Ù+8™-¸dØ$(Ø×6Ñ6ÒR¸$×:RÑ:Rð &ó 
ˆô$ 	ˆk˜9Ô%ÜÐ1Ô2Üˆk˜<×1Ñ1Ô2Ø×'Ò'ÜÐ$ l×&CÑ&CÔDØ×#Ò#Ü�(˜L×/Ñ/Ô0Ü$ \×%;Ñ%;Ó<ò 	.‰KˆAˆxØ(×/Ñ/°ÈdÐ/ÓSÐØ#×*Ñ*Ð+;Ô<Ü�Q�C�rÐ*Ð+Ð,Õ-ð	.ô
 
ˆ(„OÜ	Ð
3Ô4ä—‘˜*¬R¯X©XÔ6€JØ‡‚Ø(ò 	(ˆKØ&'ˆJ�{Ò#ð	(ð —]‘]“_×*Ñ*Ó,×3Ñ3´B·H±HÓ=Ü—h‘h §¡Ð0¼¿¹ÔAÜ—h‘h §¡Ð0¼¿¹ÔAÜ—X‘X˜tŸ~™~Ð.´b·h±hÔ?Ü "§¡¨$×*CÑ*CÐ)DÌBÏHÉHÔ UÜŸ(™( D×$7Ñ$7Ð#8ÄÇ
Á
ÔKÜ Ÿh™h¨×(?Ñ(?Ð'@ÌÏ
É
ÔSñ€Fð ‡‚Ø)ˆˆ|Ñà×!Ò!Ü#8¸ÀLÓ#QˆÐÑ à×ÒÜ˜TŸ[™[Ó)×0Ñ0×9Ñ9Ó;ˆÜ�‰�_ mÔ4Ý3à�Xˆ
Ü" :Ó.ò 	*‰IˆAˆvÜ—'‘'—,‘,˜}Ð.>ÄÀQÃÑ.GÓHˆCÙ˜S &Õ)ð	*ô ‡L�L�˜vÔ&ä$ T§[¡[°$·,±,À×@XÑ@XÓY€Kð €GÜ�4—?‘?Ó#ò ,ˆÜ—	‘	“ˆØ—‘  vÓ.ˆØ�‰”t—y‘y“{ UÑ*Õ+ð,ð —‘ Ñ#€JÝ3á ¨Ó4€Fä	ˆ.ÔØ�q‘	€IÜ	ˆ+�yÔ!Ø×#Ò#ÜÐ  &¨¡)Ô,Ø×ÒÜˆh˜˜q™	Ô"à.7¯o©oÑ+€Z� 
ØÐÜ�:Óò AˆÜ�}Ó%ò 	AˆAØ(×/Ñ/°	¸!±¸Q±ÐUYÐ/ÓZÐØ!×(Ñ(Ð)9Ô:Ü�F˜1˜#˜Z¨ s¨"Ð-=Ð,>Ð?Õ@ñ	AðAð ×ÒØ&×0Ñ0×8Ñ8¸ÀT×E^ÑE^Ð`bÓcˆÜ×(Ñ(¨Ó3ˆÜˆhŒÜÐ Ô!ÜˆoÔÜÐ%Ô&ÜˆhŒÜÐÔÜˆmÔÜÐ#Ô$ÜˆhŒà)Ð-BÑBˆÜÐ(±G©&ÀÔMØ"ˆˆxÑà×ÒÜ5ØØØØØØØó 
Ðô 	ˆoÐ3Ô4ä	ˆ%�ÔØ€Mùò[ =s   Ä 
_-c                 óæ  — t        | «      }t        |j                  «       |j                  dv rÕ|j                  rAt
        j                  j                  |j                  «      st        d|j                  › �«      ‚|j                  rAt
        j                  j                  |j                  «      st        d|j                  › �«      ‚|j                  r|j                  r|j                  r|j                  st        d«      ‚|j                  dk(  xr |j                  dk(  }|j                  dk(  r’|r�|j                  dkD  rf|j                  dk  rWt        |t        j                  «       t         j#                  d	«       |j                  d
kD  s|j$                  s|j&                  r'yt        |t        j(                  «       nt        |«       t         j#                  d«       |j                  dv rt+        ||¬«      }nt-        |||¬«      }|r`|j.                  r2t         j#                  d|j0                  › d|j0                  › d�«       |S t         j#                  d|j0                  › �«       |S )a/  Main entry function

    Args:
        argv (Optional[List[str]], optional): _description_. Defaults to None.
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Raises:
        ValueError: Path does not exist: --encoder_decoder_init_onnx
        ValueError: Path does not exist: --decoder_onnx
        ValueError: --decoder_onnx and --encoder_decoder_init_onnx are not used together for T5

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    r×  z1Path does not exist: --encoder_decoder_init_onnx z$Path does not exist: --decoder_onnx zB--decoder_onnx shall use together with --encoder_decoder_init_onnxr\   r:   ra   r_   z�The test for gpt2_sampling onnx model is limited to non-custom model with small top_p(e.g <=0.01) value. The result should be the same as gpt2 greedy search.g{®Gáz„?Nzstart testing model...)r  )r  r  zOutput files: r6   z.datazOutput file: )r   r   rD   r˜   rž   rp   rq   rà  rû   r„   r¿  rÀ  rÑ  rø  r    r0   r‹   rŒ   rÔ  r[   r/   rR  rP  rL   r›   )r2   r  r~   r  rH  s        r(   r   r   Ù  sæ  € ô  ˜4Ó €DÜ�—‘Ôà‡�˜-Ñ'Ø×)Ò)´"·'±'·.±.À×A_ÑA_Ô2`ÜÐPÐQU×QoÑQoÐPpÐqÓrÐrØ×Ò¤R§W¡W§^¡^°D×4EÑ4EÔ%FÜÐCÀD×DUÑDUÐCVÐWÓXÐXØ×*Ò*°4×3DÒ3DØ×Ò d×&DÒ&DäÐaÓbÐbà—‘ !Ñ#ÒF¨×(AÑ(AÀQÑ(F€Ià‡�˜&Ò ¡YØ�:‰:˜Ò §
¡
¨SÒ 0Ü$ T¬>×+BÑ+BÔCÜ�K‰Kð pôð �z‰z˜DÒ  D§K¢K°4·9²9Øä$ T¬>×+FÑ+FÕGä  Ô&ä
‡K�KÐ(Ô)Ø‡�˜-Ñ'Ü˜t¨yÔ9‰ä °	ÀYÔOˆáØ×(Ò(Ü�K‰K˜.¨¯©¨°R¸¿¹°}ÀEÐJÔKð €Mô �K‰K˜-¨¯© }Ð5Ô6à€Mr*   Ú__main__r%   )T)Úshared_é   NN)rV  )r   r\   rb   )NFr  )aÚ__doc__rj   ÚloggingrÆ   rp   r  Úenumr   Úpathlibr   Útypingr   r7  rÉ   r«   r  r
  r   r   Úfusion_utilsr   r	   r
   r   r  r   Útransformersr   r   r   r   r   r   r   r   Úonnxruntimer   r   r   r   Ú4onnxruntime.transformers.models.gpt2.convert_to_onnxr   r�   Ú0onnxruntime.transformers.models.gpt2.gpt2_helperr   Ú2onnxruntime.transformers.models.t5.convert_to_onnxr   rš   Ú,onnxruntime.transformers.models.t5.t5_helperr   r   Ú	getLoggerr‹   r    r·  rn   Ú	Namespacer   r�   r    Úboolr²   rÛ   rê   r  r  r   ru   Údictr<  rG  rK  r\  rc  rk  r|  r�  r«  r¼  rÊ  rú  r  rK  r_  rp  r†  rŒ  r®  r¶  r.   rø  ÚTensorr  r  rP  rR  r+   r1   r*   r(   ú<module>rh     sÒ  ðñ
#óJ Û Û Û 	Û Ý Ý Ý ã Û Û ß 4Ý $ß 4Ñ 4Ý  ÷	÷ 	ó 	÷ó õõ Tõ÷ð
 
ˆ×	Ñ	˜2Ó	€ô�Tô ñQ˜$˜s™) dÑ*ð Q°h×6HÑ6Hó Qðh*)�x×)Ñ)ó *)ðZ!�X×'Ñ'ó !ñ<O˜sð O¸dó Oñ$K¨Cð KÈ4ð KÐ[_ó Kð\ 3ð °ð ÈDð ÐUeó ðB5 §¡ð 5¸Ió 5ðpIo d§o¡oð IoÀ)ó IoðXYa°4·?±?ð YaÈyó Yað~ #ØØ$(Ø$(ñg!Øðg!àðg!ð ðg!ð ð	g!ð
 ˜T‘kðg!ð ˜T‘kóg!ðT¨:ð Àjó ð( ñØðàðð 
ˆ+Ñóò>"#òJòið.¸Jó .ðbSØ
ðSØ&*ðSØ>BðSà	óSðlf3 *ó f3ðR¨	ð Àcó ð* 9:ÐWYñ 1 	ð 1°#ð 1È4ÐPSÉ9ó 1ðhf$ Ió f$ðRD )ð DÀó DðT ØØñkØðkàðkð ðkð ð	kð
 ókð\&f¸ó &fðR\È*ó \ð~8°ó 8ñvÀ#ð Ðaeó ðF &*ñhØðhàðhð #ðhð 
ó	hòV	"0ðN '5×&?Ñ&?ñy0Ø
×
Ñ
ðy0à#óy0ðx?9Ø
×
Ñ
ð?9àÐ7Ñ7ð?9ð �|‰|ð?9ð —L‘Lð	?9ð
 ð?9ð ð?9ð ˜˜S™	‘?ð?9ð 
ˆ#ˆsˆ(�^ó?9òD	ð #'ØñTØ
×
Ñ
ðTà�C‰y˜4ÑðTð óTñnz˜×*Ñ*ð z°t¸C±yÀ4Ñ7Gó zñz8ˆt�C‰y˜4Ñð 8°4¸±9¸tÑ3Có 8ðv ˆzÒÙ…Fð r*   