o hMS@sdZddlZddlZddlZddlZddlZddlmZddlm Z ddl m Z ddl m Z mZddlmZddlmZddlmZmZdd lmZmZed ZGd d d ejZGd dde ZdddZdddZe dkruedSdS)a Prototype of ensembling N models together on the same dataset The main inference method is to run the normal transition sequence, but sum the scores for the N models and use that to choose the highest scoring transition Example of how to run it to build a silver dataset (or just parse a text file in general): # first, use this tool to build a saved ensemble python3 stanza/models/constituency/ensemble.py saved_models/constituency/wsj_inorder_?.pt --save_name saved_models/constituency/en_ensemble.pt # then use the ensemble directly as a model in constituency_parser.py python3 stanza/models/constituency_parser.py --save_name saved_models/constituency/en_ensemble.pt --mode parse_text --tokenized_file /nlp/scr/horatio/en_silver/en_split_100 --predict_file /nlp/scr/horatio/en_silver/en_split_100.inorder.mrg --retag_package en_combined_bert --lang en then, ideally, run a second time with a set of topdown models, then take the trees which match from the files N)utils)FoundationCache) BaseTrainer ModelType MultiState)Trainer)build_optimizerbuild_scheduler) ParseResult ScoredTreezstanza.constituency.trainercs eZdZd8fdd ZddZd9fdd Zed d Zed d Zed dZ ddZ ddZ ddZ eddZ eddZeddZddZddZdd Zd!d"Zd#d$Zd%d&Zd'd(Zd)d*Zd+d,Zd9d-d.Zd:d0d1Zd2d3Zd;d4d5Zd;d6d7ZZS)<EnsembleNcs.t|_|r3|rtddurtt|tr|g}tdd |fdd|D}n|s9tdt ||_ t |j D]\}}|j d|krhtd |d|||j d||j dj|jkrtd |dd ||d |dd |j djd||d |j |j dj|jkrtd|d||f|j dj|jkrtd|d||f|j d|krtd|d||f|j dj|jkrtd|d||fqD|j dj|_|tdt|j |dtj tjt|j t|jdddS)z Loads each model in filenames If foundation_cache is None, we build one on our own, as the expectation is the models will reuse modules such as pretrain, charlm, bert z6both filenames and models set when making the EnsembleNzModels used for ensemble: %s cs g|] }tj|ddjqS)F)load_optimizerfoundation_cache)rloadmodel).0filenameargsr^/var/www/html/env_mimamsha/lib/python3.10/site-packages/stanza/models/constituency/ensemble.py F z%Ensemble.__init__..z"filenames and models both not set!rz,Models {} and {} are incompatible. {} vs {}zModels z and z) are incompatible: different transitions z:  z9Models %s and %s are incompatible: different constituentsz8Models %s and %s are incompatible: different root_labelsz6Models %s and %s are incompatible: different uses_xposz=Models %s and %s are incompatible: different reverse_sentencez$Number of models in the Ensemble: %d weighted_sumT) requires_grad)super__init__r ValueErrorr isinstancestrloggerinfojoinnn ModuleListmodels enumeratetransition_schemeformat transitions constituents root_labels uses_xposreverse_sentence_reverse_sentencedetach_submodelsdebuglenregister_parametertorch Parameterzeros)selfr filenamesr(r model_idxr __class__rrr1s@   *F0zEnsemble.__init__cCs(|jD]}|D]\}}d|_q qdS)NF)r(named_parametersr)r9r_ parameterrrrr2bs zEnsemble.detach_submodelsTcs t||r|dSdSN)rtrainr2)r9moder<rrrBhs  zEnsemble.traincC |jdjSNr)r(r,r9rrrr,o zEnsemble.transitionscCrDrE)r(r.rFrrrr.srGzEnsemble.root_labelscCst|jSrA)next parametersdevicerFrrrrJwszEnsemble.devicecCstdd|jDS)zF Limit on the number of consecutive unary transitions css|]}|VqdSrA) unary_limitrmrrr sz'Ensemble.unary_limit..)minr(rFrrrrK{szEnsemble.unary_limitcC|jdSrE)r(r*rFrrrr*zEnsemble.transition_schemecCrPrE)r(has_unary_transitionsrFrrrrRrQzEnsemble.has_unary_transitionscCrDrE)r( is_top_downrFrrrrSrGzEnsemble.is_top_downcCs|jSrA)r1rFrrrr0zEnsemble.reverse_sentencecCs|jdjdS)Nr retag_method)r(rrFrrrrUszEnsemble.retag_methodcCrPrE)r(r/rFrrrr/rQzEnsemble.uses_xposcC|jd|SrE)r(get_top_constituent)r9r-rrrrWzEnsemble.get_top_constituentcCrVrE)r(get_top_transition)r9r,rrrrYrXzEnsemble.get_top_transitionc Csdg}|D]-\}}|jr4|ds4t|dk}dt|}|d|||| fqt |j D]\}}| }t |dkrT|d|||q:td|dS) NNORMS FOR MODEL PARAMETERSmodels.gư>z%.6gz %s %s %d %drz ---- MODEL %d ----r)r>r startswithr6sumabsitemnormappendnelementr)r( get_normsr4extendr#r$r%) r9linesnameparamr8r`r;rsublinesrrr log_normss  zEnsemble.log_normscCsFdg}|D]\}}|jr|d||jqtd|dS)NrZz{} {}r)r>rrar+shaper#r$r%)r9rerfrgrrr log_shapess zEnsemble.log_shapescCs0|}dd|D}|dd|jDdS)NcSs i|] \}}|ds||qS)r[)r\)rkvrrr rz'Ensemble.get_params..cSsg|]}|qSr) get_paramsrxrrrrsz'Ensemble.get_params..) base_paramschildren_params) state_dictitemsr()r9 model_staterrrros zEnsemble.get_paramscs>fdd|jD}tt|}ddt|D}|S)Ncsg|] }|qSr)initial_state_from_preterminalsrrgold_sequences gold_treespreterminal_listsrrrz.cSs g|] \}}}t|||dqS)r)rstates gold_tree gold_sequencerrrrs)r(listzip)r9r|r{rz state_batchrryrrws   z(Ensemble.initial_state_from_preterminalscpgt|D]}t|d}|durn|qtdkr6fdd|jDttddDS)z Read from the data_iterator batch_size tagged sentences and turn them into new parsing states Expects a list of list of (word, tag) Nrcg|]}|qSr)initial_state_from_wordsrxrrrrz:Ensemble.build_batch_from_tagged_words..cSg|] }t|dddqSNr~rrrrrrrr}rangerHrar4r(rr)r9 batch_size data_iteratorr?sentencerrrbuild_batch_from_tagged_wordss     z&Ensemble.build_batch_from_tagged_wordscr)zk Read from the data_iterator batch_size trees and turn them into N lists of parsing states Nrcrr)initial_state_from_gold_treesrxrrrrrz3Ensemble.build_batch_from_trees..cSrrrrrrrrr}r)r9rrr?rrrrbuild_batch_from_treess     zEnsemble.build_batch_from_treesc s`ttdd|D}ddt|j|D}tj|dd}td||j}tj|dd|}|jdtj|ddtj | ddd} fd dt t|dD}|rtt|d|D]A\}\}} | |s||ddfjd d \} } | D]} j| |rj| ||<||| f||<n qd||<d||<qf|||dfS) NcSg|]}|jqSrrrprrrrz$Ensemble.predict..cSsg|] \}}||qSr)forward)rrrrrrrr})dimz BTM,MT->BTrcsg|] }j|qSr)r,)ridxrpred_maxrrrr}T) descending)rrr(r6stackeinsumrr]argmaxtake_along_dim unsqueezedetachcpurr4r)is_legalsortr,squeeze) r9rr predictionsflat_predictionsscores pred_transrstatetransr?indicesindexrrrpredicts0    zEnsemble.predictFcsZg}ttdd|D}fddt|j|D}tt|}ddt||D}|S)NcSrrrrprrrr rz'Ensemble.bulk_apply..cs g|] \}}|j|dqS))fail) bulk_applyrrqyrr,rrr rcSsg|] \}}|j|dqS)r)_replacerrrrr )rrr()r9rr,r new_statesrrrrrs  zEnsemble.bulk_applycCsLtdt||t|}|j||j||jddd}dd|D}|S)a This parses tagged words and returns a list of trees. `parse_tagged_words` is useful at Pipeline time - it takes words & tags and processes that into trees. The tagged words should be represented: one list per sentence each sentence is a list of (word, tag) The return value is a list of ParseTree objects TODO: this really ought to be refactored with base_model zProcessing %d sentencesF) keep_statekeep_constituentscSsg|]}|jdjqS)r)rtree)rtrrrr$sz/Ensemble.parse_tagged_words..)r#r3r4evaliterparse_sentences_no_gradrr)r9wordsrsentence_iteratortreebankresultsrrrparse_tagged_wordss zEnsemble.parse_tagged_wordscsg}g} |||} ttt| } tg} |rtt} t| dkr|| \}}}|| |} tt| D]2\}}||rg| |}|j rJ| }|j }| t|t|dgdd| | ||q5tdkrfddt| D} fddt| D} t|t| D]/}t| d}|s|||}t|dkrnt|} t| d}| || t|t| qt| dks!t|| }|S)a  Repeat transitions to build a list of trees from the input batches. The data_iterator should be anything which returns the data for a parse task via next() build_batch_fn is a function that turns that data into State objects This will be called to generate batches of size batch_size until the data is exhausted The return is a list of tuples: (gold_tree, [(predicted, score) ...]) gold_tree will be left blank if the data did not include gold trees currently score is always 1.0, but the interface may be expanded to get a score from the result of the parsing transition_choice: which method of the model to use for choosing the next transition TODO: refactor with base_model rNcg|] \}}|vr|qSrr)rrrremoverrrXrz,Ensemble.parse_sentences..crrr)rr batch_idxrrrrYr)rrr4r defaultdictrsetr)finishedget_treer0reverserrar r addrHrunsort)r9rbuild_batch_fnrtransition_choicerr keep_scoresrtreebank_indicesr batch_indiceshorizon_iteratorr- pred_scoresr,rrrpredicted_treerr? horizon_state horizon_batchrrrparse_sentences'sJ             #zEnsemble.parse_sentencesc CsBt||||||||WdS1swYdSrA)r6no_gradr)r9rrrrrrrrrrrjs $z Ensemble.parse_sentences_no_grad)NNN)T)F)FFF)__name__ __module__ __qualname__rr2rBpropertyr,r.rJrKr*rRrSr0rUr/rWrYrirkrorwrrrrrrr __classcell__rrr<rr 0sB1         ! Cr cs~eZdZdZdfdd Zeddd Zd d Zed d Z ddZ eddZ eddZ eddZ edddZZS)EnsembleTrainerzj Stores a list of constituency models, useful for combining their results into one stronger model Nrr~Fc st||||||||dSrA)rr) r9ensemble optimizer schedulerepochs_trainedbatches_trainedbest_f1 best_epochfirst_optimizerr<rrrrszEnsembleTrainer.__init__cCs(t|||d}||dd}t|S)N)rrJ)r togetr)rr:rrrrr from_filesuszEnsembleTrainer.from_filescCsRg}|jjD] }|jddr!ddlm}|||j|jdq|dq|S)Nuse_peftFr)get_peft_model_state_dict) adapter_name) rr(rrpeftrra bert_model peft_name)r9paramsrrrrrget_peft_params{s   zEnsembleTrainer.get_peft_paramscCstjSrA)rENSEMBLErFrrr model_typerTzEnsembleTrainer.model_typecslfdd|jjDtfddDr#tddtdStddfd dDdS) Ncrr)num_words_knownrLrrrrrz7EnsembleTrainer.log_num_words_known..c3s|] }|dkVqdS)rNrrp)nwkrrrNsz6EnsembleTrainer.log_num_words_known..zINumber of words in the training set known to each sub-model: %d out of %drzANumber of words in the training set known to the sub-models: %srcsg|] }d|tfqS)z%d/%d)r4rprrrrr)rr(allr#r$r4r%)r9rr)rrrlog_num_words_knowns&z#EnsembleTrainer.log_num_words_knowncs,fdd}t}||_t|||}|S)Nc3s,D]\}}|ds||fVqdS)Nr[)r>r\)nprrrfake_named_parameterss   z>EnsembleTrainer.build_optimizer..fake_named_parameters)copyr>r )rrrr fake_modelrrrrr s   zEnsembleTrainer.build_optimizerc Cspt|jdj||}|dddur1z ||dW|Sty0}ztd||d}~wwtd|S)Nroptimizer_state_dictz Failed to load optimizer from %sz`Attempted to load optimizer to resume training, but optimizer not saved. Creating new optimizer) rr r(rrload_state_dictr r#r$)r checkpointrrrerrrrs zEnsembleTrainer.load_optimizercCs0t|jdj||d}d|vr||d|S)Nr)rscheduler_state_dict)r r(rr)rrrrrrrrload_schedulerszEnsembleTrainer.load_schedulerc st|tr |dn|}t|tr|dni}|dur!dgt|}|dur,dgt|}t|t|kr@tdt|t|ft|t|krTtdt|t|ffddt|||D}t|d}|j|dd |d d}|S) Nrsrrz9Model file had params length %d and peft params length %dz7Model file had params length %d and peft name length %dc s&g|]\}}}tj|||dqS))r)rmodel_from_params)r model_param peft_parampnamerrrrsz5EnsembleTrainer.model_from_params..)r(F)strictrJ) r!dictr4r rr rrr) r peft_paramsrrrrsrrr(rrrrr s"   z!EnsembleTrainer.model_from_params)NNrrr~rFrA)NN)rrr__doc__r staticmethodrrrrrr rr r rrrr<rrns"     rcCst}|jdtddd|jdtddd|jdtdddt||jdd d d |jd td ddd|jdtddddt|}|S)Nz--charlm_forward_filez$Exact path to use for forward charlm)typedefaulthelpz--charlm_backward_filez%Exact path to use for backward charlmz--wordvec_pretrain_filez'Exact name of the pretrain file to readz--langenzLanguage to use)rrr(+zWhich model(s) to load)rnargsrrz --save_nameTz#Where to save the combined ensemble)rrrequiredr)argparseArgumentParser add_argumentr"radd_device_argsvars parse_args)rparserrrrrs  rcCs6t|}t}t||d|}|j|ddddS)Nr( save_nameF)save_optimizer)rrrrsave)rrrrrrmainsr$__main__rA)!rrrloggingosr6torch.nnr&stanza.models.commonr%stanza.models.common.foundation_cacher'stanza.models.constituency.base_trainerrr stanza.models.constituency.stater"stanza.models.constituency.trainerr stanza.models.constituency.utilsr r stanza.server.parser_evalr r getLoggerr#Moduler rrr$rrrrrs0      @ ]