wenet中prefix_beam_search代码解释

转：https://blog.csdn.net/weixin_44885180/article/details/122084467

解码路径的存储形式： 每条路径prefix是一个tuple，存储对应的字索引序列如（18, 79, 1839, ...)

解码返回值：hyps:list 每个list包含一条路径和其对应的概率

　　　　　　[(prefix ,(pb, pnb)), (prefix ,(pb, pnb)) , (prefix ,(pb, pnb)) ]

    **wenet 中prefix beam search代码：**
    # cur_hyps: (prefix, (blank_ending_score, none_blank_ending_score))
    # cur_hyps (prefix ,(pb, pnb))用于当前已经存在的前缀 prefix， 初始化为 空
    #                             pb 对应前缀 prefix 以 blank 结尾的概率， 初始化为 1 （空前缀等价 blank）
    #                             pnb 对应前缀 prefix 以 非 blank 结尾的概率, 初始化为 0 （空前缀 没有 非 blank）
    cur_hyps = [(tuple(), (0.0, -float('inf')))]
    # 2. CTC beam search step by step
    for t in range(0, maxlen):
      
        logp = ctc_probs[t]  # (vocab_size,)
        # key: prefix, value (pb, pnb), default value(-inf, -inf)
        # next_hpys 保存 t 时刻，扩展的路径
        next_hyps = defaultdict(lambda: (-float('inf'), -float('inf')))
        # 2.1 First beam prune: select topk best
        # 当前 t 时刻的 beam 个 top 输出
        top_k_logp, top_k_index = logp.topk(beam_size)  # (beam_size,)
        for s in top_k_index:
            s = s.item()
            ps = logp[s].item()
            for prefix, (pb, pnb) in cur_hyps:
                last = prefix[-1] if len(prefix) > 0 else None
                if s == 0:  # blank
                    # 扩展符号为 blank, 则当前前缀扩展后还是一样，有两种情况
                    # 情况1,前缀为结尾非 blank : 如 *a (a) + _ 后 变成 *a_ (*a), 而pnb 没改变
                    # 情况2,前缀为结尾 blank : 如 *a_ (*a) + _ 后 还是为 *a_ (*a), 所以 概率更新 n_pb(*a) += (pb(*a) * ps) + (pnb((*a) * ps)
                    # 注： n_pb += 用累加是 n_pb 是时刻保存一样路径的概率和， 同理 n_pnb += 。
                    n_pb, n_pnb = next_hyps[prefix]
                    n_pb = log_add([n_pb, pb + ps, pnb + ps])
                    next_hyps[prefix] = (n_pb, n_pnb)
                elif s == last:
                    # 扩展符号非 blank 且与前缀最后字符相同，有两种情况
                    # 情况1,前缀结尾是非 blank : 如 *a + a 后 变成 (*a) , 从而更新pnb(*a)  n_pnb += pnb * ps
                    # 情况2,前缀结尾是 blank : 如 *a_ + a 后 变城 (*aa), 是另一个路径， 则新改路径 n_pnb(*aa) += pb(*a) * ps
                    #  Update *ss -> *s;
                    n_pb, n_pnb = next_hyps[prefix]
                    n_pnb = log_add([n_pnb, pnb + ps])
                    next_hyps[prefix] = (n_pb, n_pnb)
                    # Update *s-s -> *ss, - is for blank
                    n_prefix = prefix + (s, )
                    n_pb, n_pnb = next_hyps[n_prefix]
                    n_pnb = log_add([n_pnb, pb + ps])
                    next_hyps[n_prefix] = (n_pb, n_pnb)
                else:
                    # 扩展符号非 blank 且与前缀最后字符不同，则只有一种情况，就是扩展为另一个路径，此时更新这个路径概率
                    # 如 *a + c 后变后 (*ac) , *a_ + c 后也变为 (*ac) , 所以 n_pnb（*ac） += pb(*a) * ps + pnb(*a) * ps
                    n_prefix = prefix + (s, )
                    n_pb, n_pnb = next_hyps[n_prefix]
                    n_pnb = log_add([n_pnb, pb + ps, pnb + ps])
                    next_hyps[n_prefix] = (n_pb, n_pnb)
        # t 时刻重新更新保存 beam 个最优 前缀。
        next_hyps = sorted(next_hyps.items(),
                           key=lambda x: log_add(list(x[1])),
                           reverse=True)
        cur_hyps = next_hyps[:beam_size]