Spaces:
Build error
Build error
| import torch | |
| from plastic.config import ModelConfig | |
| from plastic.model.lm import PlasticDynamics, PlasticLM, build_model | |
| def test_lm_shapes_and_param_count(): | |
| cfg = ModelConfig() | |
| lm = PlasticLM(cfg) | |
| n = lm.num_params() | |
| assert 5.0e6 < n < 6.5e6, n | |
| toks = torch.randint(0, cfg.vocab_size, (2, 64)) | |
| logits, st, sig = lm(toks) | |
| assert logits.shape == (2, 64, cfg.vocab_size) and st.pos == 64 and len(sig) == cfg.n_layers | |
| def test_lm_state_carry_matches_full_pass(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, n_layers=2, chunk=16, vocab_size=50) | |
| lm = PlasticLM(cfg).to(device) | |
| toks = torch.randint(0, 50, (1, 40), device=device) | |
| full, st_full, _ = lm(toks) | |
| a, st_a, _ = lm(toks[:, :24]) | |
| b, st_b, _ = lm(toks[:, 24:], st_a) | |
| assert torch.allclose(full, torch.cat([a, b], 1), atol=1e-4, rtol=1e-4) | |
| rec, st_r, _ = lm(toks, mode="recurrent") | |
| assert torch.allclose(full, rec, atol=1e-4, rtol=1e-4) | |
| assert st_full.pos == st_b.pos == st_r.pos == 40 | |
| def test_beta_off_ablation_changes_loss(): | |
| cfg = ModelConfig(d_model=32, n_heads=2, n_layers=2, chunk=16, vocab_size=50) | |
| lm = PlasticLM(cfg) | |
| toks = torch.randint(0, 50, (2, 64)) | |
| l1 = lm.loss(toks) | |
| l0 = lm.loss(toks, beta_scale=0.0) | |
| assert not torch.allclose(l1, l0) | |
| def test_untied_head_variant(): | |
| cfg = ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=16, vocab_size=50, tie_embeddings=False) | |
| lm = PlasticLM(cfg) | |
| assert lm.head is not None | |
| logits, _, _ = lm(torch.randint(0, 50, (1, 8))) | |
| assert logits.shape == (1, 8, 50) | |
| def test_dynamics_model(): | |
| cfg = ModelConfig(domain="physics", d_model=32, n_heads=2, n_layers=2, chunk=16) | |
| m = build_model(cfg) | |
| assert isinstance(m, PlasticDynamics) | |
| x = torch.randn(2, 33, 7) | |
| pred, st, _ = m(x) | |
| assert pred.shape == (2, 33, 4) and st.pos == 33 | |
| assert m.loss(x, torch.randn(2, 33, 4)).ndim == 0 | |
| def test_domain_mismatch_rejected(): | |
| try: | |
| PlasticLM(ModelConfig(domain="physics")) | |
| except ValueError: | |
| return | |
| raise AssertionError("expected ValueError") | |