Spaces:
Sleeping
Sleeping
File size: 4,805 Bytes
dc8d52b 2d52135 dc8d52b 2d52135 dc8d52b 2d52135 dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 2d52135 dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d dc8d52b 6edd10d | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import pytest
import torch
from mini_transformer.modules.decoder import DecoderLayer, TransformerDecoder
def _dev():
return (
torch.device(f"cuda:{torch.cuda.current_device()}")
if torch.cuda.is_available()
else torch.device("cpu")
)
# -------- ctor validations --------
def test_decoder_layer_param_checks():
with pytest.raises(TypeError):
DecoderLayer("32", 4, 64, 0.1)
with pytest.raises(TypeError):
DecoderLayer(32, "4", 64, 0.1)
with pytest.raises(TypeError):
DecoderLayer(32, 4, "64", 0.1)
with pytest.raises(TypeError):
DecoderLayer(32, 4, 64, "0.1")
with pytest.raises(TypeError):
DecoderLayer(32, 4, 64, 0.1, layer_norm_style=123)
with pytest.raises(ValueError):
DecoderLayer(0, 4, 64, 0.1)
with pytest.raises(ValueError):
DecoderLayer(32, 0, 64, 0.1)
with pytest.raises(ValueError):
DecoderLayer(32, 4, 0, 0.1)
with pytest.raises(ValueError):
DecoderLayer(32, 4, 64, 1.0) # upper bound excluded
with pytest.raises(ValueError):
DecoderLayer(32, 4, 64, 0.1, layer_norm_style="weird")
def test_transformer_decoder_layers_count_checks():
with pytest.raises(TypeError):
TransformerDecoder(32, 4, 64, "2", 0.1)
with pytest.raises(ValueError):
TransformerDecoder(32, 4, 64, 0, 0.1)
# -------- forward path --------
@pytest.mark.parametrize("B,Sx,Sy,D,H,FF,L", [(2, 5, 6, 24, 3, 48, 2)])
def test_decoder_forward_happy_path(B, Sx, Sy, D, H, FF, L):
device = _dev()
dec = TransformerDecoder(D, H, FF, L, 0.1).to(device)
x = torch.randn(B, Sx, D, device=device)
y = torch.randn(B, Sy, D, device=device)
heads = dec.layers[0].self_attention_layer.num_heads
src_pad = torch.zeros(B, heads, 1, Sx, dtype=torch.bool, device=device)
tgt_pad = torch.zeros(B, heads, 1, Sy, dtype=torch.bool, device=device)
causal = (
torch.ones(B, heads, Sy, Sy, dtype=torch.bool, device=device).triu(1) if Sy > 0 else None
)
out = dec(x, y, src_pad, tgt_pad, causal)
assert out.shape == (B, Sy, D)
assert out.device == device
def test_transformer_decoder_pre_norm_layers_flag():
dec = TransformerDecoder(24, 3, 48, 2, 0.1, layer_norm_style="pre")
assert all(layer.pre_norm for layer in dec.layers)
def test_decoder_layer_pre_norm_forward_shapes():
device = _dev()
layer = DecoderLayer(24, 3, 48, 0.1, layer_norm_style="pre").to(device)
x = torch.randn(2, 5, 24, device=device)
y = torch.randn(2, 6, 24, device=device)
heads = layer.self_attention_layer.num_heads
src_mask = torch.zeros(2, heads, 1, 5, dtype=torch.bool, device=device)
tgt_mask = torch.zeros(2, heads, 1, 6, dtype=torch.bool, device=device)
causal = torch.ones(2, heads, 6, 6, dtype=torch.bool, device=device).triu(1)
out = layer(x, y, src_mask, tgt_mask, causal)
assert out.shape == y.shape
assert layer.pre_norm is True
def test_decoder_forward_input_checks_and_message_format():
layer = DecoderLayer(24, 3, 48, 0.1)
heads = layer.self_attention_layer.num_heads
base_src_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)
base_tgt_mask = torch.zeros(2, heads, 1, 3, dtype=torch.bool)
with pytest.raises(TypeError):
layer("not a tensor", torch.randn(2, 3, 24), base_src_mask, base_tgt_mask, None)
with pytest.raises(TypeError):
layer(torch.randn(2, 3, 24), "not a tensor", base_src_mask, base_tgt_mask, None)
with pytest.raises(ValueError) as e1:
layer(
torch.randn(2, 3, 24, 5),
torch.randn(2, 3, 24),
base_src_mask,
base_tgt_mask,
None,
) # x rank 4
assert "x must be a 3D torch.Tensor of shape (B, S, D)" in str(e1.value)
with pytest.raises(ValueError) as e2:
layer(
torch.randn(2, 3, 24),
torch.randn(2, 3, 24, 5),
base_src_mask,
base_tgt_mask,
None,
) # y rank 4
assert "y must be a 3D torch.Tensor of shape (B, S, D)" in str(e2.value)
x = torch.randn(2, 5, 24)
y = torch.randn(3, 6, 24) # batch mismatch
src_mask_x = torch.zeros(2, heads, 1, 5, dtype=torch.bool)
tgt_mask_y = torch.zeros(3, heads, 1, 6, dtype=torch.bool)
with pytest.raises(ValueError) as e3:
layer(x, y, src_mask_x, tgt_mask_y, None)
assert "Encoder memory and decoder input must match in batch and d_model" in str(e3.value)
y2 = torch.randn(2, 6, 16) # d_model mismatch
tgt_mask_y2 = torch.zeros(2, heads, 1, 6, dtype=torch.bool)
with pytest.raises(ValueError) as e4:
layer(x, y2, src_mask_x, tgt_mask_y2, None)
assert "Encoder memory and decoder input must match in batch and d_model" in str(e4.value)
|