o }oi4l@sdZddlZddlmZmZddlZddlmZddlm Z ddl m Z ddl m Z mZddlmZmZdd lmZdd lmZdd lmZdd lmZmZdd lmZddlmZmZddl m!Z!ddl"m#Z#dgZ$GdddeZ%dS)z This file contains code artifacts adapted from the original implementation: https://github.com/google-research/google-research/blob/master/schema_guided_dst/baseline/train_and_predict.py N)ListOptional)Trainer) DictConfig) DataLoader)DialogueSGDBERTDatasetDialogueSGDDataProcessor)evaluateget_in_domain_services)write_predictions_to_file)SGDDialogueStateLoss)NLPModel) SGDDecoder SGDEncoder) tensor2list)PretrainedModelInfo typecheck)logging)deprecated_warning SGDQAModelc s|eZdZdZeddZd4dedeffdd Ze d d Z d d Z d5de e jdededefddZd5de e jdededefddZde e jfddZd5de edefddZd5de edefddZde edede jjjdefdd Zd!d"Zd#ed$efd%d&Zd4d'eefd(d)Zd4d*eefd+d,Zd4d-eefd.d/Z dededefd0d1Z!e"dee#fd2d3Z$Z%S)6ra Dialogue State Tracking Model SGD-QA (https://arxiv.org/abs/2105.08049) The SGD-QA model is a fast multi-pass schema-guided state-tracking model, that is trained on the Google schema-guided state tracking dataset (https://arxiv.org/abs/1909.05855). The model takes dialogue as input and outputs the dialogue state, which includes slot-value pairs. The model consists of two components: a neural natural language understanding model (NLU), and a rule-based state tracker. The NLU takes in a dialogue turn and different schema (entity) information options and outputs their match score. The state tracker takes the highest rated entities and composes the dialogue state across turns. cCs|jSN)decoderselfrd/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/nlp/models/dialogue/sgdqa_model.py output_module6szSGDQAModel.output_moduleNcfgtrainercsZtdd|_tj||dt|jjj|jj j d|_ t |jjjd|_ t dd|_dS)NrF)rr) hidden_sizedropout) embedding_dimmean) reduction)r data_preparedsuper__init__r bert_modelconfigr_cfgencoderr rrr loss)rrr __class__rrr&:s zSGDQAModel.__init__c Cs^|j|||d}t|tr|d}|j|d\}}|j|||d\}}}} } } |||| | | fS)N input_idstoken_type_idsattention_maskr) hidden_states)encoded_utterancetoken_embeddingsutterance_mask)r' isinstancetupler*r) rr/r1r0r4r3logit_intent_statuslogit_req_slot_statuslogit_cat_slot_statuslogit_cat_slot_value_statuslogit_noncat_slot_status logit_spansrrrforwardDs. zSGDQAModel.forwardcCs|\}}}}}}} } } } } }}}}||||d\}}}}}}|j|||| || || || || ||d}|jjdd}|d||jd|dd||dS) Nr.r8 intent_statusr9requested_slot_statusr:categorical_slot_statusr;categorical_slot_value_statusr<noncategorical_slot_statusr=noncategorical_slot_value_startnoncategorical_slot_value_end task_maskrlr train_lossT)prog_bar)r+rH)r+ _optimizer param_groupslog)rbatch batch_idxexample_id_num service_id utterance_idsr0r1r@rArBrCrDrErFstart_char_idx end_char_idxrGr8r9r:r;r<r=r+rHrrr training_step`s^  zSGDQAModel.training_steprrNrOdataloader_idxreturncCsv|j|d\}}|d|t|jjtkr+t|jjdkr+|j|d|d|in |jd|d|id|d|iS)z Called at every validation step to aggregate and postprocess outputs on each GPU Args: batch: input batch at validation step batch_idx: batch index dataloader_idx: dataloader index rNval_losstensors) eval_step_helperrMtyperval_dataloaderslistlenvalidation_step_outputsappendrrNrOrVr+r[rrrvalidation_steps    zSGDQAModel.validation_stepcCsj|j|d\}}t|jjtkr%t|jjdkr%|j|d|d|in |jd|d|id|d|iS)z Called at every test step to aggregate and postprocess outputs on each GPU Args: batch: input batch at test step batch_idx: batch index dataloader_idx: dataloader index rXrZ test_lossr[)r\r]rtest_dataloadersr_r`test_step_outputsrbrcrrr test_steps   zSGDQAModel.test_stepc;Csp|\}}}}}}}} } } } } }}}||||d\}}}}}}|j|||||| || || || | |d}g}g}g}g}g}g}g}g}g} g}!|jjr|jjdkr|jj}"t|"D]R}#|t||t||t||t||t||t||t||t|| t||!t|qWtj ||tj ||tj ||tj ||tj ||tj ||tj ||tj ||tj | |tj |!|n2||||||||||||||||| ||!|t |}t |}t |}t |}t |}t |}t |}t |}t | }t |!}tj |}tj |}$tj j dd|}%tj|dd}&tj|%ddd}'tj |}(tj j dd|})tj|dd}*tj|)ddd}+tj j dd},|,|}-tj|-dd\}.}/|/\}0}1tj|.ddtj|/dd}2tj|1|2d ddd}3tj|1|2d ddd}4|3|4k|0dd}5t|5tj|2|2|2jd |2}2tj|2d|1ddd}6tj|2d|1dddd}7t|6|1}8t|6|1}9||||$|&|'|(|*|+|7|8|9||d }:||:fS) a. Helper called at every validation/test step to aggregate and postprocess outputs on each GPU Args: batch: input batch at step Returns: loss: averaged batch loss tensors: collection of aggregated output tensors across all GPU workers r.r?rZ)dim)axisr)device)rmdtype)rPrQr@req_slot_statuscat_slot_statuscat_slot_status_pcat_slot_value_statusnoncat_slot_statusnoncat_slot_status_p noncat_slot_pnoncat_slot_startnoncat_slot_endnoncat_alignment_startnoncat_alignment_end)r+r num_devices world_sizerangerbtorch empty_like distributed all_gathercatnnSigmoidSoftmaxargmaxmaxunbindsize unsqueezearange get_deviceviewrepeatwherezerosrn floor_dividefmod);rrNrPrQrRr0r1r@rArBrCrDrErFrSrTrGr8r9r:r;r<r=r+all_example_id_numall_service_idall_logit_intent_statusall_logit_req_slot_statusall_logit_cat_slot_statusall_logit_cat_slot_value_statusall_logit_noncat_slot_statusall_logit_spansall_start_char_idxall_end_char_idxr{indrocat_slot_status_distrprqrrnoncat_slot_status_distrsrtsoftmaxscores start_scores end_scores batch_sizemax_num_tokens total_scores start_idxend_idxinvalid_index_maskmax_span_index max_span_pspan_start_indexspan_end_indexr[rrrr\s                         zSGDQAModel.eval_step_helperoutputsc Ctdd|D}|j|dd}|j|}|j|||d}|D]\}}|j|d||ddq&|jd |ddd dS) z Called at the end of validation to post process outputs into human readable format Args: outputs: list of individual outputs of each validation step dataloader_idx: dataloader index cSg|]}|dqS)rYr.0xrrr dz9SGDQAModel.multi_validation_epoch_end..Nrirsplit dataloader_Trank_zero_onlyrYrJr)r}stackr"_validation_names_validation_dlmulti_eval_epoch_end_helperitemsrM rrrVavg_lossrrmetricskvrrrmulti_validation_epoch_end] z%SGDQAModel.multi_validation_epoch_endc Cr) z Called at the end of test to post process outputs into human readable format Args: outputs: list of individual outputs of each test step dataloader_idx: dataloader index cSr)rerrrrrrurz3SGDQAModel.multi_test_epoch_end..NrirrTrrer)r}rr" _test_names_test_dlrrrMrrrrmulti_test_epoch_endnrzSGDQAModel.multi_test_epoch_endrrc Csdtdtdtjdtfdd}dtdtfd d }td d |D}td d |D}tdd |D}tdd |D} tdd |D} tdd |D} tdd |D} tdd |D} tdd |D}tdd |D}tdd |D}tdd |D}tdd |D}tdd |D}|jjj}||j ||}i}z|j j dur|j j nd}Wnd}Y|j j dkrgt j|dd||jj j}t j|ddt|jj j||jj j}i}||d <||d!<||d"<| |d#<| |d$<| |d%<| |d&<| |d'<||d(<||d)<||d*<||d+<||d,<||d-<tt j|jj j|d.|jd/}|||jd}t||||jj|jj jd0|d1t||jj j|||jj j|jj jd2}|S)3a0 Helper called at the end of evaluation to post process outputs into human readable format Args: outputs: list of individual outputs of each step split: data split dataloader: dataloader Returns: metrics: metrics collection rids_to_service_names_dictrPrWcs fdd}tt|t|S)z Constructs string representation of example ID Args: split: evaluation data split ids_to_service_names_dict: id to service name mapping example_id_num: tensor example id c s.|\}}}}}}}d|||||||S)Nz{}-{}_{:05d}-{:02d}-{}-{}-{}-{})format) ex_id_num dialog_id_1 dialog_id_2turn_idrQ model_task_idslot_intent_idvalue_idrrrrformat_turn_idszZSGDQAModel.multi_eval_epoch_end_helper..get_str_example_id..format_turn_id)r_mapr)rrrPrrrrget_str_example_ids  zBSGDQAModel.multi_eval_epoch_end_helper..get_str_example_id predictionsrcSszddt|D}|D]-\}}|dkrt||}t|D]}|dkr.|||||<q||d|||<qq |S)a  Combines predicted values to a single example. Args: predictions: predictions ordered by keys then batch batch_size: batch size Returns: examples_preds: predictions ordered by batch then key cSsg|]}iqSrr)rrrrrrszbSGDQAModel.multi_eval_epoch_end_helper..combine_predictions_in_example.. example_idri)r|rr}chunkr)rrexamples_predsrrirrrcombine_predictions_in_examples   zNSGDQAModel.multi_eval_epoch_end_helper..combine_predictions_in_examplecSg|]}|ddqS)r[rPrrrrrrz:SGDQAModel.multi_eval_epoch_end_helper..cSr)r[rQrrrrrrrcSr)r[r@rrrrrrrcSr)r[rorrrrrrrcSr)r[rprrrrrrrcSr)r[rqrrrrrrrcSr)r[rrrrrrrrrcSr)r[rsrrrrrrrcSr)r[rtrrrrrrrcSr)r[rurrrrrrrcSr)r[rvrrrrrrrcSr)r[rwrrrrrrrcSr)r[rxrrrrrrrcSr)r[ryrrrrrrrNrzpred_res_{}_{}T)exist_okrrQr@rorprqrrrsrtrurvrwrxryz schema.jsontrainF) output_dirschemas state_tracker eval_debugin_domain_services)joint_acc_across_turnuse_fuzzy_match)strdictr}Tensorintrdialogues_processorr_services_id_to_vocabdatasetrlog_dir global_rankospathjoinrr) task_namemakedirsrget_dialogue_filesdata_dirr get_seen_servicesshaper rr rr)rrrrrrrPrQr@rorprqrrrsrtrurvrwrxryrrrprediction_dirinput_json_filesrrrrrrs     z&SGDQAModel.multi_eval_epoch_end_helpercCs8|jrdSt|jjj|jjj|j|jjd|_d|_dS)zC Preprocessed schema and dialogues and caches this N)rdialogues_example_dir tokenizerrT)r$rr)rrrrrrrrr prepare_datas zSGDQAModel.prepare_datarrcCsVtj|s t|d||jj_||jj_t d|dt d|ddS)z Update data directories Args: data_dir: path to data directory dialogues_example_dir: path to preprocessed dialogues example directory, if not exists will be created. z is not foundz"Setting model.dataset.data_dir to .z/Setting model.dataset.dialogues_example_dir to N) rrexists ValueErrorr)rrrrinfo)rrrrrrupdate_data_dirss   zSGDQAModel.update_data_dirstrain_data_configcC||j||jd|_dSN)rr)r_setup_dataloader_from_configds_item _train_dl)rrrrrsetup_training_data%zSGDQAModel.setup_training_dataval_data_configcCrr)rr r r)rrrrrsetup_validation_data)r z SGDQAModel.setup_validation_datatest_data_configcCrr)rr r r)rrrrrsetup_test_data-r zSGDQAModel.setup_test_datac Csx|jj}|j}tj|std|dt||j|jj |jj |jj |d}t j jj||j|j|j|j|j|jd}|S)Nz Data directory is not found at: r) dataset_splitrrr schema_configr)rr collate_fn drop_lastshuffle num_workers pin_memory)r)rrrrrFileNotFoundErrorrr _tokenizerrrr}utilsdatarrrrrrr)rrr dataset_cfgrrdlrrrr 1s,   z(SGDQAModel._setup_dataloader_from_configcCsg}|tdddd|S)z This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud. Returns: List of available pre-trained models. sgdqa_bertbasecasedzrhttps://api.ngc.nvidia.com/v2/models/nvidia/nemo/sgdqa_bertbasecased/versions/1.0.0/files/sgdqa_bertbasecased.nemozDialogue State Tracking model finetuned from NeMo BERT Base Cased on Google SGD dataset which has a joint goal accuracy of 59.72% on dev set and 45.85% on test set.)pretrained_model_namelocation description)rbr)clsresultrrrlist_available_modelsNsz SGDQAModel.list_available_modelsr)r)&__name__ __module__ __qualname____doc__propertyrrrr&rr>rUrr}rrrrdrhr\rrrrrrrrrrr rrr  classmethodrr% __classcell__rrr,rr+s@  "4") )&r)rtypingrrr}lightning.pytorchr omegaconfrtorch.utils.datar"nemo.collections.nlp.data.dialoguerr/nemo.collections.nlp.data.dialogue.sgd.evaluater r 7nemo.collections.nlp.data.dialogue.sgd.prediction_utilsr nemo.collections.nlp.lossesr %nemo.collections.nlp.models.nlp_modelr nemo.collections.nlp.modulesrr&nemo.collections.nlp.parts.utils_funcsrnemo.core.classes.commonrr nemo.utilsrnemo.utils.decoratorsr__all__rrrrrs&