generate.py 文件源码

def get_next_word_prob(_model, word, next_word, needModelStateReset=False):
    if needModelStateReset:
        _model.predictor.reset_state()
    _sentence_index_a = []
    index = vocab[word]
    while index != EOS_INDEX:
        y = _model.predictor(xp.array([index], dtype=xp.int32))
        probability = F.softmax(y)
        next_probs = probability.data[0]
        m = np.argsort(probability.data[0])
        break

    # In this case, the input could be an unknow word.
    if next_word not in vocab:
        return (0.0, 0.0,-1)

    next_index = vocab[next_word]
    k, = np.where(m == next_index)
    order_prob = k[0] / len(m)
    next_prob = next_probs[k[0]]

    return (order_prob, next_prob, k[0])