o }oiGl@s2ddlZddlZddlZddlmZmZddlmZddlm Z ddl m Z m Z m Z mZmZmZmZddlZddlmZddlmZddlmZmZd Zd d Zd e d e ee ffddZddZGdddeZ GdddeZ!GdddZ"de e"ded e e"fddZ#GdddeZ$Gddde$Z%dS) N)ABCabstractmethod) defaultdict)Path)AnyDictList NamedTupleOptionalTupleUnion)wer) DictConfig)TW_BREAKkaldifst_importerzriva decoder is not installed or is installed incorrectly. please run `bash scripts/installers/install_riva_decoder.sh` or `pip install riva-asrlib-decoder` to install.c Cs:zddlmmm}W|Sttfyttw)z`Import helper function that returns Riva asrlib decoder package or raises ImportError exception.rN)"riva.asrlib.decoder.python_decoderasrlibdecoderpython_decoder ImportErrorModuleNotFoundError!RIVA_DECODER_INSTALLATION_MESSAGE) riva_decoderrf/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/submodules/wfst_decoder.pyriva_decoder_importer#s rconfreturncCsFi}|D]}|ds t||}|jjdkr|nt|||<q|S)z Helper function for parsing Riva configs (namely BatchedMappedDecoderCudaConfig) into a dictionary. Args: conf: Inner Riva config. Returns: Dictionary corresponding to the Riva config. __builtins)__dir__ startswithgetattr __class__ __module___riva_config_to_dict)rresultname attributerrrr%,s   r%cCs>|D]\}}t|trtt|||qt|||qdS)a" Helper function for filling Riva configs (namely BatchedMappedDecoderCudaConfig) according to the corresponding NeMo config. Note: in-place for the first argument. Args: riva_conf: Inner Riva config. nemo_conf: Corresponding NeMo config. N)items isinstancer_fill_inner_riva_config_r"setattr) riva_conf nemo_confnemo_knemo_vrrrr+As  r+cs eZdZdZfddZZS)RivaDecoderConfigz1 NeMo config for the RivaGpuWfstDecoder. csz4t}|}d|jj_d|_d|jj_d|jj_d|j_ d|j_ d|j_ d|j_ d|jj_ t|}Wn ty?i}Ynwt|dS) N$@2g4@i'Ti )rBatchedMappedDecoderCudaConfig online_optslattice_postprocessor_optsacoustic_scalen_input_per_chunk decoder_opts default_beam max_activedeterminize_latticemax_batch_size num_channelsframe_shift_secondsword_ins_penaltyr%rsuper__init__)selfrconfigcontentr#rrrD[s"      zRivaDecoderConfig.__init__)__name__r$ __qualname____doc__rD __classcell__rrrHrr1Vsr1c@s>eZdZUdZeeed<eeed<eeed<eed<dS) WfstNbestUnitzF Container for a single RivaGpuWfstDecoder n-best hypothesis. words timesteps alignmentscoreN) rIr$rJrKr str__annotations__intfloatrrrrrMps     rMc @seZdZdZdeeeeeeeeeffddZddZ ddZ d d Z d ed e e eeeeeeeefffd dZeddZeddZeddZeddZdS)WfstNbestHypothesiszm Container for the RivaGpuWfstDecoder n-best results represented as a list of WfstNbestUnit objects. raw_hypothesescCs(t|D]Z\}}t|dtsJ|dt|dts*J|dd|dt|dts8J|dt|dtsFJ|dt|dt|dks^t|ddks^Jdqtdd|Dd d d |_t|j|_d d|jD|_t|jdj dk|_ t|jdj dk|_ dS) Nrr4z, zwords do not match timestepscSsg|]}t|qSr)rM).0rhrrr z0WfstNbestHypothesis.__init__..cS|jSN)rQ)hyprrrsz.WfstNbestHypothesis.__init__..)keycSsg|]}t|jqSr)lenrNrZhrrrr\) enumerater*tuplerUrcsorted _hypotheses_shape0_shape1rO_has_timestepsrP_has_alignment)rErWir[rrrrDs(0 zWfstNbestHypothesis.__init__ccs|jEdHdSr_rjrErrr__iter__szWfstNbestHypothesis.__iter__cCs |j|Sr_rp)rEindexrrr __getitem__s zWfstNbestHypothesis.__getitem__cCr^r_)shape0rqrrr__len__szWfstNbestHypothesis.__len__rsnew_unitcCsd|kr |jksJJ|jrt|dt|dks+|js)t|ddks+J|dkrBt|jdksq|d|j|djksq||jdkrU|j|dj|dksq|j|dj|dkrn|j|djksqJJt|tszt|}||j|<t|j|j|<dS)z Replaces a WfstNbestUnit by index. Note: in-place operation. Args: index: Index of the unit to be replaced. new_unit: Replacement unit. rr4rYN) ru has_timestepsrcrjrQr*rMrNrl)rErsrwrrr replace_unit_s &8  z!WfstNbestHypothesis.replace_unit_cCr^r_)rkrqrrrruzWfstNbestHypothesis.shape0cCr^r_)rlrqrrrshape1rzzWfstNbestHypothesis.shape1cCr^r_)rmrqrrrrxrzz!WfstNbestHypothesis.has_timestepscCr^r_)rnrqrrr has_alignmentrzz!WfstNbestHypothesis.has_alignmentN)rIr$rJrKr rRrTrUrDrrrtrvr rMrypropertyrur{rxr|rrrrrV{s&*" #   rV hypothesestokenword_disambig_strc Cs~t|}|D]}t|D]\}}g}t|jD] \}}||kr%||qt|dkrt|j} t|j} t| } t|ddkrG|| gg} } d}t|ddd|dddD]9\}}| | ||7} d | |d| t |ddd}| || | ||| |g7} |d}q]|| kr| | || 7} | | || 7} | |t | t | |j|jfq q|S)a% Searches for tokenwords in the input hypotheses and collapses them into words. Args: hypotheses: List of input WfstNbestHypothesis. tokenword_disambig_str: Tokenword disambiguation symbol (e.g. `#1`). Returns: List of WfstNbestHypothesis. rrXr4N)copydeepcopyrgrNappendrclistrOzipjoinreplacerryrhrPrQ)r~rnew_hypothesesr`kh_unit twds_listroword old_words old_timesteps words_len new_words new_timestepsj_prevjnew_wordrrrcollapse_tokenword_hypothesess>       &.    rc @seZdZdZ  d4dedededeeded ef d d Z e d5deefd dZ e defddZ e ddZ eddZejdefddZe defddZeddZejdefddZe defddZeddZejdefd dZe defd!d"Zed#d$Zed%d&Ze d'ejd(ejd)eefd*d+Ze d,eed)eefd-d.Ze d'ejd(ejd/eed)eeeffd0d1Ze d'ejd(ejd/eed)eeeeffd2d3Zd S)6AbstractWFSTDecodera Used for performing WFST decoding of the logprobs. Args: lm_fst: Language model WFST. decoding_mode: Decoding mode. E.g. `nbest`. beam_size: Beam width (float) for the WFST decoding. config: Decoder config. tokenword_disambig_id: Tokenword disambiguation index. Set to -1 to disable the tokenword mode. lm_weight: Language model weight in decoding. ?lm_fst decoding_mode beam_sizerFtokenword_disambig_id lm_weightcCsl||_||_||_|jdk|_||_d\|_|_d\|_|_d\|_ |_ |_ | || ||dS)Nr)NN)NNN)_lm_fst _beam_size_tokenword_disambig_id_open_vocabulary_decoding _lm_weight_id2word_word2id _id2token _token2id_decoding_mode_config_decoder_set_decoding_mode_set_decoder_config _init_decoder)rErrrrFrrrrrrDs       zAbstractWFSTDecoder.__init__NcCdSr_rrErFrrrr)z'AbstractWFSTDecoder._set_decoder_configcCrr_rrErrrrr-rz&AbstractWFSTDecoder._set_decoding_modecCrr_rrqrrrr1rz!AbstractWFSTDecoder._init_decodercCr^r_)rrqrrrr5rzz!AbstractWFSTDecoder.decoding_modevaluecC||dSr_)_decoding_mode_setterrErrrrr9cCrr_rrrrrr=rz)AbstractWFSTDecoder._decoding_mode_settercCr^r_)rrqrrrrArzzAbstractWFSTDecoder.beam_sizecCrr_)_beam_size_setterrrrrrErcCrr_rrrrrrIrz%AbstractWFSTDecoder._beam_size_settercCr^r_)rrqrrrrMrzzAbstractWFSTDecoder.lm_weightcCrr_)_lm_weight_setterrrrrrQrcCrr_rrrrrrUrz%AbstractWFSTDecoder._lm_weight_settercCr^r_)rrqrrrrYrzz)AbstractWFSTDecoder.tokenword_disambig_idcCr^r_)rrqrrropen_vocabulary_decoding]rzz,AbstractWFSTDecoder.open_vocabulary_decoding log_probslog_probs_lengthrcCdS) Decodes logprobs into recognition hypotheses. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. Returns: List of recognition hypotheses. Nr)rErrrrrdecodeaszAbstractWFSTDecoder.decoder~cCr) Does various post-processing of the recognition hypotheses. Args: hypotheses: List of recognition hypotheses. Returns: List of processed recognition hypotheses. NrrEr~rrr _post_decoders z AbstractWFSTDecoder._post_decodereference_textscCr) Calibrates LM weight to achieve the best WER for given logprob-text pairs. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. reference_texts: List of reference word sequences. Returns: Pair of (best_lm_weight, best_wer). NrrErrrrrrcalibrate_lm_weightz'AbstractWFSTDecoder.calibrate_lm_weightcCr) Calculates the oracle (the best possible WER for given logprob-text pairs. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. reference_texts: List of reference word sequences. Returns: Pair of (oracle_wer, oracle_wer_per_utterance). Nrrrrrcalculate_oracle_werrz(AbstractWFSTDecoder.calculate_oracle_wer)rrr_) rIr$rJrKrrRrUr rTrDrrrrr}rsetterrrrrrrrtorchTensorrrrr rrrrrrrs           rcseZdZdZ      d=ded eefd ed ed ed de dede ffdd Z d>d ed fddZ ddZ d efddZ defddZdefddZdefddZed d!Zejdefd"d!Zdefd#d$Zd%ejd&ejd'eefd(d)Zd%ejd&ejd'eefd*d+Zd%ejd&ejd'ed,fd-d.Zd%ejd&ejd'eeeed,ffd/d0Zd1eeeed,fd'eeeed,ffd2d3Zd%ejd&ejd4eed'eeeffd5d6Zd%ejd&ejd4eed'eeeeffd7d8Z d9d:Z!d;d<Z"Z#S)?RivaGpuWfstDecoderaU Used for performing WFST decoding of the logprobs with the Riva WFST decoder. Args: lm_fst: Kaldi-type language model WFST or its path. decoding_mode: Decoding mode. Choices: `nbest`, `mbr`, `lattice`. beam_size: Beam width (float) for the WFST decoding. config: Riva Decoder config. tokenword_disambig_id: Tokenword disambiguation index. Set to -1 to disable the tokenword mode. lm_weight: Language model weight in decoding. nbest_size: N-best size for decoding_mode == `nbest` mbrr2Nrrr4rzkaldifst.StdFstrrrFr1rr nbest_sizecs&||_d|_t||||||dSr_) _nbest_size_load_word_latticerCrD)rErrrrFrrrrHrrrDs zRivaGpuWfstDecoder.__init__cCsj|dus t|dkr t}t|dsttd|j|jj_|j |jj j |jj _ |j |jj _||_dS)Nrr7z/Unexpected config error. Please debug manually.)rcr1hasattrr RuntimeErrorrr7r; lattice_beamrr8r9lm_scalernbestrrrrrrs    z&RivaGpuWfstDecoder._set_decoder_configcst}t}ddlm}||_|j}d}d}t|ttfr$|j |}n$t||j |j fr?|}t j dd}||j|j}n tdt||jdd}|jdurdd t|jd D|_|jtt|j}|d tfd d |_|D] \} } | |j| <q|jdurdd t|jd D|_|jtt|j} | d tfdd |_| D] \} } | |j| <qt j dd#} |j | j|!} t"| |j#|$| || j||_%Wdn1swY|r|&dSdS)Nr)load_word_latticezw+t)modezUnsupported lm_fst type: z#0r4cS*i|]}t|dd|ddqS r4rrTsplitrZlinerrr z4RivaGpuWfstDecoder._init_decoder.. zcSr_rr) word_unk_idrrraz2RivaGpuWfstDecoder._init_decoder..cSrrrrrrrr rcrr_rr) token_unk_idrrrar)'rr+nemo.collections.asr.parts.utils.wfst_utilsrrrr*rrR StdVectorFstread StdConstFsttempfileNamedTemporaryFilewriter' ValueErrortype input_symbolsfindroutput_symbolsstriprr#mapreversedr)rrrr write_textr6r+rBatchedMappedDecoderCudarclose)rEkaldifstrrrtmp_fst tmp_fst_filenum_tokens_with_blankword2idrvtoken2id words_tmprFr)rrrrsZ          z RivaGpuWfstDecoder._init_decodercCsN|dkr |j|_n|dkr|j|_n|dkr|j|_ntd|||_dS)NrrlatticezUnsupported mode: ) _decode_nbest_decode _decode_mbr_decode_latticerrrrrrrs    z%RivaGpuWfstDecoder._set_decoding_modercC4|j|kr|||jjj_|||_dSdSr_)r_release_gpu_memoryrr7r;rrrrrrr*   z$RivaGpuWfstDecoder._beam_size_settercCs@|j|kr|||jjjj|jjj_|||_dSdSr_)rrrr7r8r9rrrrrrr1s   z$RivaGpuWfstDecoder._lm_weight_settercCs|j|kr ||dSdSr_)rrrrrrr:s z(RivaGpuWfstDecoder._decoding_mode_settercCr^r_)rrqrrrr>rzzRivaGpuWfstDecoder.nbest_sizecCrr_)_nbest_size_setterrrrrrBrcCrr_)rrrr7r8rrrrrrrFrz%RivaGpuWfstDecoder._nbest_size_setterrrrc Cs|j||}g}|D]R}g}|D]B}gg}} t|j|jD]\} } | dkr6||j| | t| qdd|jD} |j } |t t |t | t | | gq|t t |q |S)a Decodes logprobs into recognition hypotheses via the N-best decoding decoding. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. Returns: List of WfstNbestHypothesis with empty alignment and trivial score. rcSsg|]}|dqS)r4r)rZilabelrrrr\gr]z4RivaGpuWfstDecoder._decode_nbest..) r decode_nbestrrNword_start_times_secondsrrrTilabelsrQrhrV)rErrhypotheses_nbestr~nhnbest_containerrerNrOwtrPrQrrrrMs  $z RivaGpuWfstDecoder._decode_nbestc Cs~|j||}g}|D]1}gg}}|D]}||d|t|dq|tttt|t|tdggq |S)a Decodes logprobs into recognition hypotheses via the Minimum Bayes Risk (MBR) decoding. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. Returns: List of WfstNbestHypothesis with empty alignment and trivial score. rr4r5)r decode_mbrrrTrVrh) rErrhypotheses_mbrr~rerNrOerrrrms ,zRivaGpuWfstDecoder._decode_mbrKaldiWordLatticec st:}|jd}|j||ddtt|Dd||||j|jfddtt|D}Wd|S1sAwY|S)ax Decodes logprobs into kaldi-type lattices. Args: log_probs: A torch.Tensor of the predicted log-probabilities of shape [Batch, Time, Vocabulary]. log_probs_length: A torch.Tensor of length `Batch` which contains the lengths of the log_probs elements. Returns: List of KaldiWordLattice. z.latscSsg|]}t|qSrrRrZrorrrr\r]z6RivaGpuWfstDecoder._decode_lattice..zark,t:csg|]}t|qSrrrhypotheses_latticerrr\sN) rrr'rdecode_write_latticerangercrr)rErrtmp_lat tmp_lat_namer~rrrrs       z"RivaGpuWfstDecoder._decode_latticecCs8|}|tjd}|||}||}|S)rcpu) contiguoustorlongrr)rErrr~rrrrs   zRivaGpuWfstDecoder.decoder~cCs&|jr|jdvrt||j|jS|S)r)rr)rrrrrrrrrrs zRivaGpuWfstDecoder._post_decoderc Cst|t|ks J|j}|j}d|_dtd}}tddD]&}|d|_|||} tdd| D|} t|| | |krE|j| }}q||_||_||fS) rrginfr4 cSsg|] }d|djqS) r)rrNrdrrrr\sz:RivaGpuWfstDecoder.calibrate_lm_weight..)rcrrrUrrword_error_rateprint) rErrrdecoding_mode_backuplm_weight_backupbest_lm_weightbest_werrr~r rrrrs     z&RivaGpuWfstDecoder.calibrate_lm_weightc sjrtt|t|ksJj}d_||}ggg}}}t||D]/\} } fdd| D} || rAt| nd|| | ||d|dq(|_t |t ||fS)rrcsg|]}j|qSr)r)rZrrqrrr\rfz;RivaGpuWfstDecoder.calculate_oracle_wer..r4r) rNotImplementedErrorrcrrrrrr edit_distancesum) rErrrr'latticesscorescounts wer_per_uttrtextword_idsrrqrrs z'RivaGpuWfstDecoder.calculate_oracle_wercCs(z|`Wn ty YnwtdS)zS Forces freeing of GPU memory by deleting the Riva decoder object. N)r Exceptiongccollectrqrrrr s   z&RivaGpuWfstDecoder._release_gpu_memorycCs |dSr_)rrqrrr__del__s zRivaGpuWfstDecoder.__del__)rr2Nrrr4r_)$rIr$rJrKr rrRrUr rTrDrrrrrrr}rrrrrrrVrrrrrr rrrr7rLrrrHrrs 9         # " r)&rr5rabcrr collectionsrpathlibrtypingrrrr r r r rjiwerr r% omegaconfrrrrrrrRr%r+r1rMrVrrrrrrrs:  $    P 03