decode.py 文件源码

def simple_search(lattice, queues, rnn_predictor, beam_size):
    # Simple but slow implementation of beam search
    for i in range(len(lattice)):
        for j in range(len(lattice[i])):
            for target, word_id in lattice[i][j]:
                for previous_cost, previous_string, previous_state, previous_prediction in queues[j]:
                    cost = previous_cost + previous_prediction[word_id]
                    string = previous_string + target
                    predictions, states = rnn_predictor.predict([word_id], [previous_state])
                    hypothesis = (cost, string, states[0], predictions[0])
                    queues[i].append(hypothesis)

        # prune queues[i] to beam size
        queues[i] = heapq.nsmallest(beam_size, queues[i], key=operator.itemgetter(0))
    return queues