Simple attempt at beam search.

beam_search.py

import numpy as np
import heapq

VOCAB_SIZE = 1000
HIDDEN_DIM = 128

vocab = {
    'the': 5,
    'fox': 35,
    'jumped': 144,
}

W_embed = np.random.randn(VOCAB_SIZE * HIDDEN_DIM).reshape(
    VOCAB_SIZE, HIDDEN_DIM)
W_i = np.random.randn(HIDDEN_DIM * HIDDEN_DIM).reshape(
    HIDDEN_DIM, HIDDEN_DIM)
W_h = np.random.random(HIDDEN_DIM * HIDDEN_DIM).reshape(
    HIDDEN_DIM, HIDDEN_DIM)
W_o = np.random.random(VOCAB_SIZE * HIDDEN_DIM).reshape(
    VOCAB_SIZE, HIDDEN_DIM)


def relu(x):
    switch = ((np.random.random(x.shape) - .5) > 0.0).astype(np.float)
    return x * switch


def softmax(x):
    max_x = max(x)
    exp_x = np.exp(x - max_x)
    normalizer = exp_x.sum()
    return exp_x / normalizer


def rnn_fun(input_, hidden_):
    hidden_ = np.tanh(np.dot(W_i, input_) + np.dot(W_h, hidden_))
    output_ = np.log(softmax(np.dot(W_o, hidden_)) + 1e-9)
    return hidden_, output_


def score_sequence(words):
    word_ids = [vocab[word] for word in words]
    word_vectors = W_embed[word_ids, :]

    top_indices = []

    hidden_ = np.zeros(HIDDEN_DIM, dtype=np.float)
    for word_id, word_vector in zip(word_ids, word_vectors):
        hidden_, output_ = rnn_fun(word_vector, hidden_)
        highest_index = np.argmax(output_)
        top_indices.append((output_[highest_index], highest_index))

    return top_indices


def beam_one(current_candidate, top_k=10):
    accumulated_score, list_of_word_ids, hidden_ = current_candidate

    word_id = list_of_word_ids[-1]
    word_vector = W_embed[word_id, :]
    hidden_, output_ = rnn_fun(word_vector, hidden_)

    top_indices = np.argsort(output_)[::-1][:top_k]  # highest first

    new_candidates = []
    for top_index in top_indices:
        new_score = output_[top_index]
        new_candidates.append((
            accumulated_score + new_score,
            list_of_word_ids + [top_index],
            hidden_)
        )

    return new_candidates


def beam_search(seed_word):
    word_id = vocab[seed_word]
    hidden_ = np.zeros(HIDDEN_DIM, dtype=np.float)

    visited = set()
    candidates = [(0.0, [word_id], hidden_)]

    seq_length = 1
    while seq_length < 10:
        new_candidates = []
        for current_candidate in candidates:
            new_candidates.extend(beam_one(current_candidate))

        for new_candidate in new_candidates:
            if tuple(new_candidate[1]) in visited:
                continue
            else:
                visited.add(tuple(new_candidate[1]))

            heapq.heappush(candidates, new_candidate)

        while len(candidates) > 10:
            garbage = heapq.heappop(candidates)
            assert garbage[0] < candidates[0][0]

        seq_length += 1

    return candidates


def _test():
    print(score_sequence('the fox jumped'.split()))
    candidates = beam_search('the')
    for candidate in candidates:
        print(candidate[0], candidate[1])


if __name__ == '__main__':
    _test()
#     [(-2.0054691832436857, 544), (-1.7653964179241393, 469), (-1.2171965458736256, 813)]
#     (-3.292619036283746, [5, 553])
#     (-3.1967043001394249, [5, 164])
#     (-3.1523385115177573, [5, 256])
#     (-3.0905989819888835, [5, 172, 960])
#     (-2.8036062169513656, [5, 172])
#     (-2.7562216474461989, [5, 808])
#     (-2.651661954532639, [5, 964])
#     (-2.0054691832436857, [5, 544])
#     (-2.3370670925794239, [5, 440])
#     (0.0, [5])