test_encoder_decoder.py

'''Unit tests for encoder_decoder.py'''


import torch
import encoder_decoder


def test_update_beam():
    torch.manual_seed(1030)
    N, K, V, H = 2, 2, 5, 10
    ed = encoder_decoder.EncoderDecoder(
        encoder_decoder.Encoder, encoder_decoder.DecoderWithAttention,
        V, V,
        encoder_hidden_size=H,
        cell_type='rnn', beam_width=2,
    )
    logpb_tm1 = torch.arange(K).flip(0).unsqueeze(0).expand(N, -1).float()
    logpb_tm1 -= 1.5
    logpb_tm1[1] *= 2  # [[-0.5, -1.5], [-1., -3.]]
    htilde_t = torch.rand(N, K, 2 * H)
    logpy_t = (
        torch.arange(V).unsqueeze(0).unsqueeze(0)
        .expand(N, K, -1).float() * -1.1
    )  # [x, y, :] = [0., -1.1, -2.2, ...]
    # [0, x, :] = [0, 1]
    b_tm1_1 = torch.arange(K).unsqueeze(0).unsqueeze(0).expand(-1, N, -1)
    b_t_0, b_t_1, logpb_t = ed.update_beam(
        htilde_t, b_tm1_1, logpb_tm1, logpy_t)
    # batch 0 picks path 0 extended with 0, then path 1 extended with 0
    assert torch.allclose(logpb_t[0], torch.tensor([-0.5, -1.5]))
    assert torch.allclose(b_t_0[0, 0], htilde_t[0, 0])
    assert torch.allclose(b_t_0[0, 1], htilde_t[0, 1])
    assert torch.allclose(b_t_1[:, 0, 0], torch.tensor([0, 0]))
    assert torch.allclose(b_t_1[:, 0, 1], torch.tensor([1, 0]))
    # batch 0 picks path 0 extended with 0, then path 0 extended with 1
    assert torch.allclose(logpb_t[1], torch.tensor([-1., -2.1]))
    assert torch.allclose(b_t_0[1, 0], htilde_t[1, 0])
    assert torch.allclose(b_t_0[1, 1], htilde_t[1, 0])
    assert torch.allclose(b_t_1[:, 1, 0], torch.tensor([0, 0]))
    assert torch.allclose(b_t_1[:, 1, 1], torch.tensor([0, 1]))