decode_ngram.py 文件源码

def search(lattice, ngrams, queues, beam_size, viterbi_size):
    for i in range(len(lattice)):
        for j in range(len(lattice[i])):
            for target, source in lattice[i][j]:

                word_queue = []
                for previous_cost, previous_history in queues[j]:
                    history = previous_history + [(target, source)]
                    cost = previous_cost + get_ngram_cost(ngrams, tuple(history[-3:]))
                    hypothesis = (cost, history)
                    word_queue.append(hypothesis)

                # prune word_queue to viterbi size
                if viterbi_size > 0:
                    word_queue = heapq.nsmallest(viterbi_size, word_queue, key=operator.itemgetter(0))

                queues[i] += word_queue

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