| """Regression tests for the P2 SpaceByte hierarchy. |
| |
| These tests pin the patch-boundary contract separately from the full model and |
| prove the hierarchy can exactly collapse to the P1 plain decoder when every byte |
| position is promoted to a global patch. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from ksbyte.config import BOS_ID, ByteLMConfig |
| from ksbyte.model import ByteDecoder, SpaceByteDecoder, build_model, build_spacebyte_boundary_mask |
|
|
|
|
| def _tiny_cfg(**overrides) -> ByteLMConfig: |
| cfg = ByteLMConfig( |
| d_model=32, |
| n_layers=2, |
| n_heads=4, |
| n_kv_heads=2, |
| mlp_ratio=2.0, |
| ctx_len=16, |
| dropout=0.0, |
| n_local_in=0, |
| n_global=2, |
| n_local_out=0, |
| max_patches=16, |
| tie_embeddings=True, |
| ) |
| return cfg.merge(overrides).validate() |
|
|
|
|
| def test_spacebyte_boundary_mask_marks_first_spacelike_byte_only(): |
| |
| x = torch.tensor([[BOS_ID, ord("a"), 32, 32, ord("b"), 9, ord("c")]]) |
| seg_ids = torch.zeros_like(x) |
|
|
| mask = build_spacebyte_boundary_mask(x, seg_ids=seg_ids) |
|
|
| assert mask.tolist() == [[True, False, True, False, False, True, False]] |
|
|
|
|
| def test_spacebyte_gather_scatter_reuses_latest_patch_context(): |
| x = torch.tensor([[BOS_ID, ord("a"), 32, ord("b"), ord("c"), 32, ord("d")]]) |
| h = torch.arange(x.numel() * 3, dtype=torch.float32).view(1, x.numel(), 3) |
| boundary = build_spacebyte_boundary_mask(x) |
|
|
| gathered, patch_mask, patch_pos, patch_seg, patch_ids = SpaceByteDecoder.gather_patches( |
| h, |
| boundary, |
| pos_ids=torch.arange(x.size(1)).unsqueeze(0), |
| seg_ids=torch.zeros_like(x), |
| max_patches=8, |
| ) |
| scattered = SpaceByteDecoder.scatter_patches(gathered + 1000.0, patch_ids) |
|
|
| assert patch_mask.tolist() == [[True, True, True]] |
| assert patch_pos.tolist() == [[0, 2, 5]] |
| assert patch_seg.tolist() == [[0, 0, 0]] |
| |
| assert torch.equal(scattered[0, 0], gathered[0, 0] + 1000.0) |
| assert torch.equal(scattered[0, 1], gathered[0, 0] + 1000.0) |
| assert torch.equal(scattered[0, 2], gathered[0, 1] + 1000.0) |
| assert torch.equal(scattered[0, 4], gathered[0, 1] + 1000.0) |
| assert torch.equal(scattered[0, 5], gathered[0, 2] + 1000.0) |
| assert torch.equal(scattered[0, 6], gathered[0, 2] + 1000.0) |
|
|
|
|
| def test_spacebyte_factory_builds_model(): |
| model = build_model(_tiny_cfg(variant="spacebyte")) |
| assert isinstance(model, SpaceByteDecoder) |
|
|
|
|
| def test_spacebyte_reduces_to_plain_decoder_when_all_positions_are_boundaries(): |
| torch.manual_seed(7) |
| plain_cfg = _tiny_cfg(variant="plain") |
| space_cfg = _tiny_cfg(variant="spacebyte") |
| plain = ByteDecoder(plain_cfg).eval() |
| space = SpaceByteDecoder(space_cfg).eval() |
|
|
| mapped = {} |
| for name, tensor in plain.state_dict().items(): |
| mapped[name.replace("blocks.", "global_blocks.")] = tensor |
| missing, unexpected = space.load_state_dict(mapped, strict=False) |
| assert not unexpected |
| assert set(missing) == set() |
|
|
| x = torch.randint(0, 128, (2, 8)) |
| pos_ids = torch.arange(x.size(1)).unsqueeze(0).expand_as(x) |
| all_boundaries = torch.ones_like(x, dtype=torch.bool) |
|
|
| plain_logits, _, _ = plain(x, pos_ids=pos_ids) |
| space_logits, _, _ = space(x, pos_ids=pos_ids, boundary_mask=all_boundaries) |
|
|
| torch.testing.assert_close(space_logits, plain_logits, atol=1e-5, rtol=1e-5) |
|
|
