o i|!@s~dZddlZddlZddlZddlmZddlZddlZde de dfddZ dd d Z dd d Z GdddeZdS)zCommon functions for ASR.N)groupbyic st|dkrdSd}t|ddddd}tj|D].}||fdd|D}t|dkrIt|d dddd}|d |d |krI|d 7}q||krPdSdS) zEnd detection. described in Eq. (50) of S. Watanabe et al "Hybrid CTC/Attention Architecture for End-to-End Speech Recognition" :param ended_hyps: :param i: :param M: :param D_end: :return: rFcS|dSNscorexrrN/home/ubuntu/.local/lib/python3.10/site-packages/espnet/nets/e2e_asr_common.py!zend_detect..T)keyreversecs g|] }t|dkr|qS)yseq)len.0r  hyp_lengthrr %s zend_detect..cSrrrr rrr r (r rr)rsortedsixmovesrange) ended_hypsiMD_endcountbest_hypmhyps_same_lengthbest_hyp_same_lengthrrr end_detects&   r$c Cs|dur t|d}t|d}Wdn1swY|dkr{|dus.Jd|t|}|D]$\}}tdd|dd d D} t| d kr[|| d 7<q7t|||d <d ||d k<d ||<| tj t |} | St d |t| S) zObtain label distribution for loss smoothing. :param odim: :param lsm_type: :param blank: :param transcript: :return: Nrbuttsunigramz-transcript is required for %s label smoothingcSsg|]}t|qSr)int)rnrrr rGz(label_smoothing_dist..outputrtokenidrz*Error: unexpected label smoothing type: %s)openjsonloadnpzerositemsarraysplitrastypefloat32sumloggingerrorsysexit) odimlsm_type transcriptblankf trans_json labelcountkvids labeldistrrr label_smoothing_dist4s,   $  rGcCsH||}ttj|tjdd}ttj|tjdd}t||S)zReturn the output size of the VGG frontend. :param in_channel: input channel size :param out_channel: output channel size :return: output size :rtype int )dtype)r0ceilr3r6r()idim in_channel out_channelrrr get_vgg2l_odimVs rOcsNeZdZdZ dfdd ZdddZddZd d Zd d Zd dZ Z S)ErrorCalculatoraCalculate CER and WER for E2E_ASR and CTC models during training. :param y_hats: numpy array with predicted text :param y_pads: numpy array with true (target) text :param char_list: :param sym_space: :param sym_blank: :return: Fcsftt|||_||_||_||_||_|j|j|_ |j|jvr.|j|j|_ dSd|_ dS)z$Construct an ErrorCalculator object.N) superrP__init__ report_cer report_wer char_listspacer?index idx_blank idx_space)selfrU sym_space sym_blankrSrT __class__rr rRos  zErrorCalculator.__init__cCshd\}}|r |||S|js|js||fS|||\}}|jr'|||}|jr0|||}||fS)adCalculate sentence-level WER/CER score. :param torch.Tensor ys_hat: prediction (batch, seqlen) :param torch.Tensor ys_pad: reference (batch, seqlen) :param bool is_ctc: calculate CER score for CTC :return: sentence-level WER score :rtype float :return: sentence-level CER score :rtype float )NN)calculate_cer_ctcrSrTconvert_to_char calculate_cer calculate_wer)rZys_hatys_padis_ctccerwerseqs_hat seqs_truerrr __call__s     zErrorCalculator.__call__cCs0ddl}gg}}t|D]x\}}ddt|D}||} gg} } |D]} t| } | dkrC| |jkrC| |jkrC| |jt| q%| D]} t| } | dkrd| |jkrd| |jkrd| |jt| qFd| } d| }t |dkr|| | ||t |q |rt t |t |}|Sd}|S)zCalculate sentence-level CER score for CTC. :param torch.Tensor ys_hat: prediction (batch, seqlen) :param torch.Tensor ys_pad: reference (batch, seqlen) :return: average sentence-level CER score :rtype float rNcSsg|]}|dqS)rrrrrr rr*z5ErrorCalculator.calculate_cer_ctc..) editdistance enumeraterr(rXrYappendrUjoinrevalfloatr7)rZrcrdrmcers char_ref_lensryy_haty_trueseq_hatseq_trueidx hyp_chars ref_charscer_ctcrrr r_s4     z!ErrorCalculator.calculate_cer_ctccsgg}}t|D]Z\}}||}t|dkd}t|dkr$|dnt|} fdd|d| D} fdd|D} d| jd} | jd} d| jd} || || q ||fS) a"Convert index to character. :param torch.Tensor seqs_hat: prediction (batch, seqlen) :param torch.Tensor seqs_true: reference (batch, seqlen) :return: token list of prediction :rtype list :return: token list of reference :rtype list rkrcsg|] }jt|qSr)rUr(rrzrZrr rsz3ErrorCalculator.convert_to_char..Ncs&g|]}t|dkrjt|qS)rk)r(rUr~rrr rs&rl ) rnr0whererrpreplacerVr?ro)rZrcrdrhrirrvrweos_trueymaxrxry seq_hat_text seq_true_textrrr r`s   zErrorCalculator.convert_to_charc Csxddl}gg}}t|D]$\}}||}|dd} |dd} ||| | |t| q tt|t|S)zCalculate sentence-level CER score. :param list seqs_hat: prediction :param list seqs_true: reference :return: average sentence-level CER score :rtype float rNrrl)rmrnrrorqrrrr7) rZrhrirmchar_edsrtrrrr{r|rrr ras   zErrorCalculator.calculate_cerc Cspddl}gg}}t|D] \}}||}|} |} ||| | |t| q tt|t|S)zCalculate sentence-level WER score. :param list seqs_hat: prediction :param list seqs_true: reference :return: average sentence-level WER score :rtype float rN)rmrnr4rorqrrrr7) rZrhrirmword_eds word_ref_lensrrr hyp_words ref_wordsrrr rbs zErrorCalculator.calculate_wer)FF)F) __name__ __module__ __qualname____doc__rRrjr_r`rarb __classcell__rrr]r rPds  "rP)Nr)rrH)rr.r8r: itertoolsrnumpyr0rlogexpr$rGrOobjectrPrrrr s  " "