beam_search.py 文件源码

def _calc_top_n(self, model, x, state, beam_width):
        o, state = model.decode_once(x, state, train=False)
        o = F.log_softmax(o, use_cudnn=False)
        o = chainer.cuda.to_cpu(o.data[0])

        eos_score = o[self.EOS]
        self._edit_probs(o)

        inds = np.argpartition(o, len(o) - beam_width)
        inds = inds[::-1][:beam_width]
        return inds, o, state, eos_score