o i+@sHdZddlmZmZddlZddlmZddlmZGddde Z dS)z'Error Calculator module for Transducer.)ListTupleN) AbsDecoder)BeamSearchTransducercseZdZdZ  ddedejjdee de de de d e ffd d Z d ej d ej fddZdej d ej deeeffddZdej dej defddZdej dej defddZZS)ErrorCalculatorTransduceraCalculate CER and WER for transducer models. Args: decoder: Decoder module. token_list: List of tokens. sym_space: Space symbol. sym_blank: Blank symbol. report_cer: Whether to compute CER. report_wer: Whether to compute WER. Fdecoder joint_network token_list sym_space sym_blank report_cer report_wercsFtt||dddd|_||_||_||_||_||_||_ dS)z'Construct an ErrorCalculatorTransducer.defaultF)rr beam_size search_type score_normN) super__init__r beam_searchrr spaceblankr r )selfrrr r r r r  __class__[/home/ubuntu/.local/lib/python3.10/site-packages/espnet2/asr/transducer/error_calculator.pyrs  z"ErrorCalculatorTransducer.__init__ encoder_outtargetc Csd\}}t|d}g}|t|jj}t|D]}|||}| |qdd|D} | | |\} } |j rC| | | }|j rL|| | }||fS)a.Calculate sentence-level WER/CER score for Transducer model. Args: encoder_out: Encoder output sequences. (B, T, D_enc) target: Target label ID sequences. (B, L) Returns: : Sentence-level CER score. : Sentence-level WER score. )NNrcSsg|] }|djddqS)rN)yseq).0 nbest_hyprrr Msz6ErrorCalculatorTransducer.__call__..)intsizetonextr parametersdevicerangerappendconvert_to_charr calculate_cerr calculate_wer) rrrcerwer batchsize batch_nbestb nbest_hypspred char_pred char_targetrrr__call__6s     z"ErrorCalculatorTransducer.__call__r5returnc sgg}}t|D]D\}}fdd|D}fdd||D}d|jd}|jd}d|jd}|jd}||||q ||fS)aEConvert label ID sequences to character sequences. Args: pred: Prediction label ID sequences. (B, U) target: Target label ID sequences. (B, L) Returns: char_pred: Prediction character sequences. (B, ?) char_target: Target character sequences. (B, ?) cg|] }jt|qSrr r$)r!hrrrr#jz=ErrorCalculatorTransducer.convert_to_char..cr:rr;)r!rr=rrr#kr> ) enumeratejoinreplacerrr+) rr5rr6r7ipred_i char_pred_i char_target_irr=rr,Ys   z)ErrorCalculatorTransducer.convert_to_charr6r7c Cstddl}gg}}t|D]"\}}|dd}||dd} |||| |t| q tt|t|S)zCalculate sentence-level CER score. Args: char_pred: Prediction character sequences. (B, ?) char_target: Target character sequences. (B, ?) Returns: : Average sentence-level CER score. rNrAr@) editdistancerBrDr+evallenfloatsum rr6r7rI distanceslensrErGr5rrrrr-xs  z'ErrorCalculatorTransducer.calculate_cerc Cslddl}gg}}t|D]\}}|}||} |||| |t| q tt|t|S)zCalculate sentence-level WER score. Args: char_pred: Prediction character sequences. (B, ?) char_target: Target character sequences. (B, ?) Returns: : Average sentence-level WER score rN)rIrBsplitr+rJrKrLrMrNrrrr.s  z'ErrorCalculatorTransducer.calculate_wer)FF)__name__ __module__ __qualname____doc__rtorchnnModulerr$strboolrTensorr8rr,rLr-r. __classcell__rrrrr sR#   r) rUtypingrrrVespnet2.asr.decoder.abs_decoderr-espnet2.asr.transducer.beam_search_transducerrobjectrrrrrs