详解Wenet中的Prefix Beam Search算法实现
最编程
2024-07-24 14:44:42
...
转：https://blog.****.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]
原文地址：https://www.cnblogs.com/wieneralan/p/16906221.html