o i@sdZddlmZddlZddlmZddlmZddlmZddlmZddlm Z dd lm Z ddl Z dd l m Z dd lmZdd lmZGd ddeZGddde jjZGddde jjZdS)zBeam search module.)chainN)Any)Dict)List) NamedTuple)Tuple)Union) end_detect)PartialScorerInterface)ScorerInterfacec@sveZdZUdZejed<dZee ejfed<e Z e e ee ejffed<e Ze e efed<de fdd Zd S) HypothesiszHypothesis data type.yseqrscorescoresstatesreturncCs0|j|jt|jdd|jDdS)z#Convert data to JSON-friendly dict.cSsi|] \}}|t|qS)float).0kvrrS/home/ubuntu/.local/lib/python3.10/site-packages/funasr/models/scama/beam_search.py sz%Hypothesis.asdict..)r rr)_replacer tolistrrritems_asdictselfrrrasdictszHypothesis.asdictN)__name__ __module__ __qualname____doc__torchTensor__annotations__rrrdictrrstrrrrrrrrr s   r csBeZdZdZ   dde j d1e j de j d2ed3ed4ed5ede e fd6d7Zd8ed4ed2ed,e e d9e e de e f d:d;ZZS)?BeamSearchScamaBeam search implementation.N?scorersweights beam_size vocab_sizesoseos token_listpre_beam_ratiopre_beam_score_keyc Rt||_t|_t|_t|_tj |_ | D]E\} } | | d} | dks0| dur1qt | tsBJ| dt| d| |j| <t | trR| |j| <n| |j| <t | tjjrc| |j | <q||_||_||_t|||_||_||_| dur| dkr| |jvrt| d|j| |_|jduo|j|jkot|jdk|_dSaTInitialize beam search. Args: scorers (dict[str, ScorerInterface]): Dict of decoder modules e.g., Decoder, CTCPrefixScorer, LM The scorer will be ignored if it is `None` weights (dict[str, float]): Dict of weights for each scorers The scorer will be ignored if its weight is 0 beam_size (int): The number of hypotheses kept during search vocab_size (int): The number of vocabulary sos (int): Start of sequence id eos (int): End of sequence id token_list (list[str]): List of tokens for debug log pre_beam_score_key (str): key of scores to perform pre-beam search pre_beam_ratio (float): beam size in the pre-beam search will be `int(pre_beam_ratio * beam_size)` rNz (z$) does not implement ScorerInterfacefullz is not found in super__init__r-r'r, full_scorers part_scorersr$nn ModuleDictnn_dictrget isinstancer typer Moduler0r1r2int pre_beam_sizer.n_vocabKeyErrorr4len do_pre_beam rr,r-r.r/r0r1r2r3r4rrw __class__rrr:'L            zBeamSearchScama.__init__xrcCXt}t}|jD]\}}||||<d||<q td||tj|jg|jddgSzGet an initial hypothesis data. Args: x (torch.Tensor): The encoder output feature Returns: Hypothesis: The initial hypothesis. device)rrrr r'r,r init_stater r$tensorr0rTrrO init_states init_scoresrdrrrinit_hypq  zBeamSearchScama.init_hypxscC$tj|g|j|jd}t||fSzAppend new token to prefix tokens. Args: xs (torch.Tensor): The prefix token x (int): The new token to append Returns: torch.Tensor: New tensor contains: xs + [x] with xs.dtype and xs.device dtyperTr$rWrbrTcatr^rOrrr append_token zBeamSearchScama.append_tokenhypx_maskpre_acoustic_embedsc CsPt}t}|jD]\}}|j|j|j||||d\||<||<q ||fS)Score new hypothesis by `self.full_scorers`. Args: hyp (Hypothesis): Hypothesis with prefix tokens to score x (torch.Tensor): Corresponding input feature Returns: Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of score dict of `hyp` that has string keys of `self.full_scorers` and tensor score values of shape: `(self.n_vocab,)`, and state dict that has string keys and state values of `self.full_scorers` rirjr'r;rrr r) rrhrOrirjrrrr[rrr score_fullszBeamSearchScama.score_fullidscCLt}t}|jD]\}}||j||j||\||<||<q ||fSaaScore new hypothesis by `self.part_scorers`. Args: hyp (Hypothesis): Hypothesis with prefix tokens to score ids (torch.Tensor): 1D tensor of new partial tokens to score x (torch.Tensor): Corresponding input feature Returns: Tuple[Dict[str, torch.Tensor], Dict[str, Any]]: Tuple of score dict of `hyp` that has string keys of `self.part_scorers` and tensor score values of shape: `(len(ids),)`, and state dict that has string keys and state values of `self.part_scorers` r'r<r score_partialr rrrhrorOrrrr[rrrrs &zBeamSearchScama.score_partialweighted_scorescCz|d|dkr||jd}||fS||}td |dd<|||<||jd}|||jd}||fSaCompute topk full token ids and partial token ids. Args: weighted_scores (torch.Tensor): The weighted sum scores for each tokens. Its shape is `(self.n_vocab,)`. ids (torch.Tensor): The partial token ids to compute topk Returns: Tuple[torch.Tensor, torch.Tensor]: The topk full token ids and partial token ids. Their shapes are `(self.beam_size,)` rinfNsizetopkr.rrrvrotop_idstmp local_idsrrrbeamzBeamSearchScama.beam prev_scoresnext_full_scoresfull_idxnext_part_scorespart_idxcCVt}|D]\}}||||||<q|D]\}}||||||<q|SaMerge scores for new hypothesis. Args: prev_scores (Dict[str, float]): The previous hypothesis scores by `self.scorers` next_full_scores (Dict[str, torch.Tensor]): scores by `self.full_scorers` full_idx (int): The next token id for `next_full_scores` next_part_scores (Dict[str, torch.Tensor]): scores of partial tokens by `self.part_scorers` part_idx (int): The new token id for `next_part_scores` Returns: Dict[str, torch.Tensor]: The new score dict. Its keys are names of `self.full_scorers` and `self.part_scorers`. Its values are scalar tensors by the scorers. r'rrrrrr new_scoresrrrrr merge_scores zBeamSearchScama.merge_scoresr part_statescCLt}|D]\}}|||<q|jD]\}}||||||<q|SaMerge states for new hypothesis. Args: states: states of `self.full_scorers` part_states: states of `self.part_scorers` part_idx (int): The new token id for `part_scores` Returns: Dict[str, torch.Tensor]: The new score dict. Its keys are names of `self.full_scorers` and `self.part_scorers`. Its values are states of the scorers. r'rr< select_staterrrr new_statesrrr[rrr merge_states   zBeamSearchScama.merge_states running_hypscCsfg}tj|j|jd}|D]}tj|j|j|jd}|j||||d\} } |jD] } ||j| | | 7}q(|j rN|j dkr@|n| |j } t | |j d}| |||\} }|jD]} |||j| | | 7<qZ||j7}t|||D]#\}}|t||||j|||j| || ||| ||dqyt|ddd d d tt||j}q |S) "Search new tokens for running hypotheses and encoded speech x. Args: running_hyps (List[Hypothesis]): Running hypotheses on beam x (torch.Tensor): Encoded speech feature (T, D) Returns: List[Hypotheses]: Best sorted hypotheses rSrarlr7ryrr rrcS|jSNrrOrrrUz(BeamSearchScama.search..TkeyreverseNr$arangerFrTzerosrbrnr;r-rIr4r}rErsr<rziprappendr rfr rrrsortedminrHr.)rrrOrirj best_hypspart_idsrhrvrrrpre_beam_scores part_scoresrjpart_jrrrsearchs@        zBeamSearchScama.searchrRr scama_mask maxlenratio minlenratiomaxlenminlenc sf|dur.|dkr|jd}n|dkrdt|}n tdt||d}t||d}tdt|jdtdt|tdt||}g} t|D]} t dt| d} |dur|d} t | | d} |dd| | dddf} d}|dur|\}}}t j |d|f|j d j|jd }t j||fdd }t | |} |dd| | dddf}j||| |d }| |||| }|d krtdd| D| rtd| nt|dkrtdn t dt|qWt| dddd}t|dkr'td|dkrgS||td |dS|D]}dfdd|jD}t d|j||jq)|d}|jD] \}}t|ddj|dd|j|dd |qOtd!|jd"td#|jt|jd"td$t|jdurtd%dfd&d|jddDd'|S)(WPerform beam search. Args: x (torch.Tensor): Encoded speech feature (T, D) maxlenratio (float): Input length ratio to obtain max output length. If maxlenratio=0.0 (default), it uses a end-detect function to automatically find maximum hypothesis lengths If maxlenratio<0.0, its absolute value is interpreted as a constant max output length. minlenratio (float): Input length ratio to obtain min output length. Returns: list[Hypothesis]: N-best decoding results Nrrydecoder input length: max output length: min output length: position rbrSdimrlrRcSg|]}|qSrrrhrrr z+BeamSearchScama.forward..end detected at no hypothesis. Finish decoding.remained hypotheses: cSrrrrrrrrrz)BeamSearchScama.forward..TrOthere is no N-best results, perform recognition again with smaller minlenratio.皙?cg|]}j|qSrr2rrOrrrr!nbest: y: {}, yseq: {}, score: {}6.2f * 3 =  for total log probability: .2fnormalized log probability: "total number of ended hypotheses: best hypo: crrrrrrrrr ) shaperDmaxr|logginginfor(r\rangedebugrr$rrbtorTrdr post_processr rHrwarningforwardjoinr formatrrrr-r2)rrOrrjrrrrr ended_hypsimask_enctoken_num_predictortoken_id_slicepre_acoustic_embeds_curbtr[padbest nbest_hypsr rrrrrrZs        : (zBeamSearchScama.forwardrrc  tdt|jdur'tddfdd|djddD||dkr;td fd d|D}g}|D]D}|jd jkr~tj j D]#\}} | |j |} |j|| 7<|j|jj|| d }qT||q?||q?|S) Perform post-processing of beam search iterations. Args: i (int): The length of hypothesis tokens. maxlen (int): The maximum length of tokens in beam search. maxlenratio (int): The maximum length ratio in beam search. running_hyps (List[Hypothesis]): The running hypotheses in beam search. ended_hyps (List[Hypothesis]): The ended hypotheses in beam search. Returns: List[Hypothesis]: The new running hypotheses. "the number of running hypotheses: Nrrcrrrrrrrrrz0BeamSearchScama.post_process..rry-adding in the last position in the loopc$g|]}|j|jjdqS)r rrfr r1rrrrrrrrrrHr2rr rr1rr;rr< final_scorerrrrr-r rrrrrr remained_hypsrhrr[srrrr( (     zBeamSearchScama.post_processNr+N)NN)NNrRrRNr)r r!r"r#rr(r rrDrr:r$r%r r\ staticmethodrfrrrnrsrrrrrr __classcell__rrrLrr)$s     J         > kr)csNeZdZdZ   d=deeefdeeefdededed ed e ed ed effd d Z de e fddZ e dejdedejfddZddifde dejdejdejdedeeeejfeeefff ddZde dejdejdeeeejfeeefffddZdejdejdeejejffd d!Ze d"eeefd#eeejfd$ed%eeejfd&edeeejff d'd(Zd)ed*ed&edefd+d,Zddifd-e e dejdejdejdede e f d.d/Zddd0d0dd1ifdejd2ejdejd3ed4ed5ed6edede e fd7d8Zd9ed5ed3ed-e e d:e e de e f d;d<ZZS)>BeamSearchScamaStreamingr*Nr+r,r-r.r/r0r1r2r3r4c r5r6r8rJrLrrr:rNz!BeamSearchScamaStreaming.__init__rcCrPrQrUrXrrrr\Cr]z!BeamSearchScamaStreaming.init_hypr^rOcCr_r`rcrerrrrf[rgz%BeamSearchScamaStreaming.append_tokenrhrirjcachec CsRt}t}|jD]\}} | j|j|j|||||d\||<||<q ||fS)rkrirjrrm) rrhrOrirjrrrrr[rrrrnjsz#BeamSearchScamaStreaming.score_fullrocCrprqrrrtrrrrsruz&BeamSearchScamaStreaming.score_partialrvcCrwrxr{r~rrrrrzBeamSearchScamaStreaming.beamrrrrrcCrrrrrrrrrz%BeamSearchScamaStreaming.merge_scoresrrcCrrrrrrrrrz%BeamSearchScamaStreaming.merge_statesrcCshg}tj|j|jd}|D]}tj|j|j|jd} |j|||||d\} } |jD] } | |j| | | 7} q)|j rO|j dkrA| n| |j } t | |j d}| |||\}}|jD]} | ||j| || 7<q[| |j7} t|| |D]#\}}|t| |||j|||j| ||||| ||dqzt|ddd d d tt||j}q |S) rrSrarr7ryrcSrrrrrrrr.rz1BeamSearchScamaStreaming.search..TrNr)rrrOrirjrrrrhrvrrrrrrrrrrrrs@         zBeamSearchScamaStreaming.searchrRrrrrrrc  s*|dur.|dkr|jd}n|dkrdt|}n tdt||d}t||d}tdt|jdtdt|tdt||d} g} t|D]} td t| d} d} |dur|\}}}t j |d|f|j d j |j d }t j||fdd }t| |}|dd||dddf} j| || | |d d}| |||| } |dkrtdd| D| rtd| nt| dkrtdn tdt| qVt| dddd}t|dkr td|dkrgS||td|dS|D]}dfdd|jD}td|j||jq |d}|jD] \}}t|ddj|d d!|j|dd"|q1td#|jd$td%|jt|jd$td&t|jdurtd'dfd(d|jddDd)|S)*rNrrryrrrrrrrSrdecoderrrRcSrrrrrrrrrz4BeamSearchScamaStreaming.forward..rrrcSrrrrrrrrrz2BeamSearchScamaStreaming.forward..Trrrrcrrrrrrrrrrrrrrrrrrrrcrrrrrrrrrr)rrDrr|rrr(rrr$rrbrrTrdrrrr rHrrrrr rrrrr-r2)rrOrrjrrrrrrrrrrrrr[rrrrr rrrrrr3s     : (z BeamSearchScamaStreaming.forwardrrc r) rrNrrcrrrrrrrrrz9BeamSearchScamaStreaming.post_process..rryrcrrrrrrrrrrrrrrrrrrz%BeamSearchScamaStreaming.post_processr)r r!r"r#rr(r rrDrr:r r\rr$r%rfr'rrrnrsrrrrrrrrrrLrrs     J #        ?  qr)r# itertoolsrrtypingrrrrrrr$funasr.metrics.commonr 2funasr.models.transformer.scorers.scorer_interfacer r r r=rCr)rrrrrs&          U