diff options
Diffstat (limited to 'tests/test_data.py')
| -rw-r--r-- | tests/test_data.py | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/tests/test_data.py b/tests/test_data.py index fe84e94..cd4e620 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -30,7 +30,7 @@ def test_encode_decode_roundtrip_integers(): cfg = Config(vocab_mode="integers", range_end=100) for n in [2, 42, 100]: assert decode_tokens(encode_int(n, cfg) + [cfg.eos_id], cfg) == n - assert cfg.vocab == 102 # 0..101 + EOS + assert cfg.vocab == 103 # values 0..101, EOS=102, pad=103 (no alias with target 101) def test_splits_sizes_and_no_overlap(): @@ -60,8 +60,8 @@ def test_make_batch_shapes_and_padding(): cfg = Config() ex = build_examples([2, 7, 42, 99], cfg) b = make_batch(ex, cfg) - assert b["x"].shape == (4, 2) # max input len 2 (42, 99) - assert b["y"].shape == (4, 4) # max output len 4 ("101"+EOS) + assert b["x"].shape == (4, 3) # GLOBAL layout: left-padded to 3 digits (max of [2,100]) + assert b["y"].shape == (4, 4) # GLOBAL layout: right-padded to 4 ("101"+EOS) assert b["y_in"].shape == (4, 4) assert b["y_mask"].shape == (4, 4) # pad positions masked out @@ -69,6 +69,9 @@ def test_make_batch_shapes_and_padding(): assert b["y_mask"][0, 1] # BOS position of y_in is eos_id assert (b["y_in"][:, 0] == cfg.eos_id).all() + # LEFT-padding: input 2 sits at the last column, pads at the front + assert b["x"][0, 2] == 2 and b["x"][0, 0] == cfg.pad_id + assert b["x"][2, 1] == 4 and b["x"][2, 2] == 2 def test_flagged_composites_are_composite(): |