File size: 1,159 Bytes
19b8775 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import pytest
import torch
from stanza import Pipeline
from stanza.models.constituency.positional_encoding import SinusoidalEncoding, AddSinusoidalEncoding
from stanza.tests import *
pytestmark = [pytest.mark.pipeline, pytest.mark.travis]
def test_positional_encoding():
encoding = SinusoidalEncoding(model_dim=10, max_len=6)
foo = encoding(torch.tensor([5]))
assert foo.shape == (1, 10)
# TODO: check the values
def test_resize():
encoding = SinusoidalEncoding(model_dim=10, max_len=3)
foo = encoding(torch.tensor([5]))
assert foo.shape == (1, 10)
def test_arange():
encoding = SinusoidalEncoding(model_dim=10, max_len=2)
foo = encoding(torch.arange(4))
assert foo.shape == (4, 10)
assert encoding.max_len() == 4
def test_add():
encoding = AddSinusoidalEncoding(d_model=10, max_len=4)
x = torch.zeros(1, 4, 10)
y = encoding(x)
r = torch.randn(1, 4, 10)
r2 = encoding(r)
assert torch.allclose(r2 - r, y, atol=1e-07)
r = torch.randn(2, 4, 10)
r2 = encoding(r)
assert torch.allclose(r2[0] - r[0], y, atol=1e-07)
assert torch.allclose(r2[1] - r[1], y, atol=1e-07)
|