gTYDddlZddlZddlZddlmZddlmZmZmZddl Z ddl m Z m Z mZddlmZde deded ed e j$f d Zd e deded e j$fd ZdedefdZ d-de dededededed e j$fdZdedeee j$ffdZdededed ededede d e de dededefdZdededededede d e de dededefdZdZ d.ded eed!eed"eed eee j$ee j$ee j$ff d#Z d.d$ed eed!eed"eed eee j$ee j$ee j$ff d%Zd&Zd'ed(edededededed eed!eed"eed)edededefd*Z d+Z!e"d,k(re!yy)/N)Path)DictOptionalTuple) ModelProto TensorProto numpy_helper) OnnxModel input_ids batch_sizesequence_lengthdictionary_sizereturnc(|jjjtjtj tj fvsJtjj|||ftj}|jjjtjk(rtj|}|S|jjjtj k(rtj|}|S)a`Create input tensor based on the graph input of input_ids Args: input_ids (TensorProto): graph input of the input_ids input tensor batch_size (int): batch size sequence_length (int): sequence length dictionary_size (int): vocabulary size of dictionary Returns: np.ndarray: the input tensor created )sizedtype) type tensor_type elem_typerFLOATINT32INT64nprandomrandintint32float32int64)r r r rdatas \/var/www/openai/venv/lib/python3.12/site-packages/onnxruntime/transformers/bert_test_data.pyfake_input_ids_datar!s >> % % / /4   99  _J3PXZX`X`  aD~~!!++{/@/@@zz$ K  # # - -1B1B Bxx~ K segment_idsc|jjjtjtj tj fvsJtj||ftj}|jjjtjk(rtj|}|S|jjjtj k(rtj|}|S)a,Create input tensor based on the graph input of segment_ids Args: segment_ids (TensorProto): graph input of the token_type_ids input tensor batch_size (int): batch size sequence_length (int): sequence length Returns: np.ndarray: the input tensor created r) rrrrrrrrzerosrrr)r#r r rs r fake_segment_ids_datar'2s    ' ' 1 16   88Z1 BD##--1B1BBzz$ K    % % / /;3D3D Dxx~ Kr"max_sequence_lengthaverage_sequence_lengthc|dk\r||ksJd|z|kDrtjd|z|z |Stjdd|zdz S)N)rr)r(r)s r get_random_lengthr-Msd "a ',CGZ,ZZ Z ""%88~~a"99*QRUVQVtJNQ,>'?$q1u+'M\]Z!^$YZ]^Y^tJNQ4F/G$qSTu+/UdeQ^a'!+,+:&1Q':>*'(+B'BZ!^$/03J/JQ^a'!+,+"",, 0A0AAzz$ K  $ $ . .+2C2C Cxx~ Kr" directoryinputsc  tjj|s& tj|t d|dnt d|dt |jD]t\}\}}tj||}ttjj|d|dd 5}|j|jd d d vy #t $rt d|dYwxYw#1swYxYw) zOutput input tensors of test data to a directory Args: directory (str): path of a directory inputs (Dict[str, np.ndarray]): map from input name to value z#Successfully created the directory  zCreation of the directory z failedzWarning: directory z$ existed. Files will be overwritten.input_.pbwbN)ospathexistsmkdirprintOSError enumerateitemsr from_arrayopenjoinwriteSerializeToString)r;r<indexnamertensorfiles r output_test_datarSs 77>>) $ F HHY  7 {!D E #I;.RST(8|d((t4 "'',,yF5'*=> E JJv//1 2F E 9 C .ykA B CF EsC&; D&DDD test_casesverbose random_seedc |Jtjj|tj|g} t|D]} t ||||}|j |i}|rt |||||j <|rt|||| | | ||j <|rt| dk(r td|| j|| S)aCreate given number of input data for testing Args: batch_size (int): batch size sequence_length (int): sequence length test_cases (int): number of test cases dictionary_size (int): vocabulary size of dictionary for input_ids verbose (bool): print more information or not random_seed (int): random seed input_ids (TensorProto): graph input of input IDs segment_ids (TensorProto): graph input of token type IDs input_mask (TensorProto): graph input of attention mask average_sequence_length (int): average sequence length excluding paddings random_sequence_length (bool): whether use uniform random number for sequence length mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key Returns: List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary with input name as key and a tensor as value rzExample inputs) rrseedr4r!rPr'r:lenrFappend)r r rTrrUrVr r#r.r)r/r0 all_inputs _test_caseinput_1r<s r fake_test_datar^sD   IINN; KK JJ' %i_o^..'* '<[*Ve'fF;## $ &:J9PRhjs'F:?? # s:!+ "F +&!( r"rXc hd} t|||| ||||||| | } t| |k7r td| S)aCreate given number of input data for testing Args: batch_size (int): batch size sequence_length (int): sequence length test_cases (int): number of test cases seed (int): random seed verbose (bool): print more information or not input_ids (TensorProto): graph input of input IDs segment_ids (TensorProto): graph input of token type IDs input_mask (TensorProto): graph input of attention mask average_sequence_length (int): average sequence length excluding paddings random_sequence_length (bool): whether use uniform random number for sequence length mask_type (int): mask type 1 is mask index; 2 is 2D mask; 3 is key len, cumulated lengths of query and key Returns: List[Dict[str,numpy.ndarray]]: list of test cases, where each test case is a dictionary with input name as key and a tensor as value i'z$Failed to create test data for test.)r^rYrF) r r rTrXrUr r#r.r)r/r0rr[s r generate_test_datar`sU@O  J :*$ 45 r"c|t|jk\ry|j|}|j|}|A|j||}|-|jdk(r|j|jd}|S)NCastr)rYinputfind_graph_input get_parentop_type) onnx_model embed_node input_indexrc graph_input parent_nodes r get_graph_input_from_embed_noderl%sc***++   [ )E--e4K ++J D  "{':':f'D$55k6G6G6JKK r"rginput_ids_namesegment_ids_nameinput_mask_namecp|j}||j|}|td|d}|r!|j|}|td|d}|r!|j|}|td|d|rdndz|rdndz}t||k7rtd|dt||||fSt|dk7rtdt||j d } t| dk(rh| d} t || d}t || d}t || d }|(|D]#} | j j} d | vs"| }%| td |||fSd}d}d}|D]0} | j j} d | vr| }$d | vsd| vr| }/| }2|r |r|r|||fStd)a Find graph inputs for BERT model. First, we will deduce inputs from EmbedLayerNormalization node. If not found, we will guess the meaning of graph inputs based on naming. Args: onnx_model (OnnxModel): onnx model object input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None. segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None. input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None. Raises: ValueError: Graph does not have input named of input_ids_name or segment_ids_name or input_mask_name ValueError: Expected graph input number does not match with specified input_ids_name, segment_ids_name and input_mask_name Returns: Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids, segment_ids and input_mask Nz Graph does not have input named r+rzExpect the graph to have z inputs. Got r2z'Expect the graph to have 3 inputs. Got EmbedLayerNormalizationmaskz#Failed to find attention mask inputtokensegmentz?Fail to assign 3 inputs. You might try rename the graph inputs.)'get_graph_inputs_excluding_initializersrd ValueErrorrYget_nodes_by_op_typerlrPlower) rgrmrnro graph_inputsr r#r.expected_inputs embed_nodesrhrcinput_name_lowers r find_bert_inputsr~2s>4EEGL!//?  ??OPQ Q $556FGK" #CDTCU!VWW #44_EJ! #COCT!UVVKqQ7 1PQR |  /88IWZ[gWhVijk k+z11 <AB3|CTBUVWW112KLK ;1 ^ 3J AN 5j*aP 4ZQO  %#(::#3#3#5 --!&J&  BC C+z11IKJ ::++- % %J ' '98H+HKI[Z+z11 V WWr" onnx_filect}t|d5}|j|jdddt |}t ||||S#1swY"xYw)aFind graph inputs for BERT model. First, we will deduce inputs from EmbedLayerNormalization node. If not found, we will guess the meaning of graph inputs based on naming. Args: onnx_file (str): onnx model path input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None. segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None. input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None. Returns: Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids, segment_ids and input_mask rbN)rrKParseFromStringreadr r~)rrmrnromodelrRrgs r get_bert_inputsrsW( LE i $ diik* 5!J J8H/ ZZ  s AA!cttj}|jddtd|jddtdd|jd dtd d |jd dtd d|jddtdd|jddtdd|jddtdd|jddtd d|jddtdd|jdddd|j d|jdddd |j d!|jd"d#d$td%&|jd'd(ddd)|j d*|jd+dtd,d-|j }|S).Nz--modelTzbert onnx model path.)requiredrhelpz --output_dirFz4output test data path. Default is current directory.)rrdefaultrz --batch_sizer+zbatch size of inputz--sequence_lengthz maximum sequence length of inputz--input_ids_namezinput name for input idsz--segment_ids_namezinput name for segment idsz--input_mask_namezinput name for attention maskz --samplesz$number of test cases to be generatedz--seedr2z random seedz --verbose store_truezprint verbose information)ractionr)rUz--only_input_tensorsz-only save input tensors and no output tensors)only_input_tensorsz-az--average_sequence_lengthz)average sequence length excluding padding)rrrz-rz--random_sequence_lengthz3use uniform random instead of fixed sequence length)r/z --mask_typer,z^mask type: (1: mask index, 2: raw 2D mask, 3: key lengths, cumulated lengths of query and key))argparseArgumentParser add_argumentstrint set_defaults parse_args)parserargss r parse_argumentsrs'  $ $ &F  DsAXY   C  S!Rgh   /    '    )    ,    3  5sAMZ  (   &  <   51  #  8   " B  u5   m     D Kr"r output_dirrc t|||| \}}}t||||||||| | | }t|D]=\}}tjj |dt |z}t||?| ryddl}d|jvrddgndg}|j||}|jDcgc]}|j}}t|D]\}}tjj |dt |z}|j||}t|D]\}}tjt!j"|||}t%tjj |d|dd 5}|j'|j)dddycc}w#1swYxYw) aICreate test data for a model, and save test data to a directory. Args: model (str): path of ONNX bert model output_dir (str): output directory batch_size (int): batch size sequence_length (int): sequence length test_cases (int): number of test cases seed (int): random seed verbose (bool): whether print more information input_ids_name (str): graph input name of input_ids segment_ids_name (str): graph input name of segment_ids input_mask_name (str): graph input name of input_mask only_input_tensors (bool): only save input tensors, average_sequence_length (int): average sequence length excluding paddings random_sequence_length (bool): whether use uniform random number for sequence length mask_type(int): mask type test_data_set_NrCUDAExecutionProviderCPUExecutionProvider) providersoutput_r@rA)rr`rHrBrCrLrrS onnxruntimeget_available_providersInferenceSession get_outputsrPrunr rJrasarrayrKrMrN)rrr r rTrXrUrmrnrorr)r/r0r r#r.r[r6r<r;rrsessionoutput output_namesresult output_name tensor_resultrRs r create_and_save_test_datarsD*9P`bq)r&I{J#  Jz* 6GGLL-=A-FG F++ #k&I&I&K K !"89$ % **5I*FG.5.A.A.CD.CFFKK.CLDz* 6GGLL-=A-FG \62' 5NA{(33BJJvay4I;WMbggll9s#.>?F$ =::<=GF6+EGFs8F4 F99G ct}|jdkr|j|_|j}|Yt |j }t jj|jd|jd|j}|t |}|jddn tdt|j ||j|j|j|j|j |j"|j$|j&|j(|j|j*|j,td|y)Nrbatch__seq_T)parentsexist_okz7Directory existed. test data files will be overwritten.z Test data is saved to directory:)rr)r rrrrBrCrLparentr rErFrsamplesrXrUrmrnrorr/r0)rrprCs r mainrZs  D ##q('+';';$J  WW\\!((fT__4EU4K_K_J`,ab J 4$ / GH           $$ ## " ,j9r"__main__)r,)NNN)#rrBrpathlibrtypingrrrnumpyronnxrrr rgr rndarrayr!r'r-boolr:rrSr^r`rlr~rrrr__name__r"r rs ((66 (+>ATWZZ<{VY^`^h^h6B3BB  FFFF! F ! F  FZZFR33T#rzz/-B3.7777 7  7  7777!7!77t1111  1  1  111!1!11h %)&*%) YXYXSMYXsmYXc] YX  8BJJ "**!5x 7K KL YX|%)&*%) [[SM[sm[c] [  8BJJ "**!5x 7K KL [8aHI> I>I>I> I>  I>  I>I>SMI>smI>c]I>I>!I>!I>I>X$:N zFr"