tests/test_checkpoint.py

1import json23import numpy as np4import pytest56from checkpoint import save_checkpoint, load_checkpoint7from tokenizer import build_tokenizer8from transformer import ModelConfig, init_params, forward910CORPUS = "low low low lower newest newest\nthe low newest tower\n"111213@pytest.fixture14def saved(tmp_path):15    tok = build_tokenizer(CORPUS, max_merges=8)16    cfg = ModelConfig(tok.vocab_size, context_len=8, emb_dim=4,17                      num_heads=2, ff_dim=8, num_layers=2)18    params = init_params(cfg, np.random.default_rng(0))19    prefix = str(tmp_path / "ckpt" / "model")        # "ckpt" does not exist yet20    save_checkpoint(prefix, params, cfg, tok)21    return prefix, params, cfg, tok222324def test_save_creates_the_three_files(saved):25    prefix = saved[0]26    for ext in (".npz", ".json", ".tok.json"):27        assert open(prefix + ext, "rb").read()       # exists and is not empty282930def test_config_round_trips(saved):31    prefix, _, cfg, _ = saved32    assert load_checkpoint(prefix)[1] == cfg333435def test_weights_round_trip_exactly(saved):36    prefix, params, _, _ = saved37    loaded = load_checkpoint(prefix)[0]38    assert loaded.keys() == params.keys()39    for k in params:40        assert np.array_equal(loaded[k], params[k]), k414243def test_tokenizer_round_trips(saved):44    prefix, _, _, tok = saved45    t2 = load_checkpoint(prefix)[2]46    assert t2.itos == tok.itos47    assert t2.stoi == tok.stoi48    assert t2.vocab_size == tok.vocab_size49    for s in ["low newest", "lower tower", "  low  ", "low!"]:   # "!" was never seen50        assert t2.encode(s) == tok.encode(s), s51        assert t2.decode(t2.encode(s)) == tok.decode(tok.encode(s)), s52    assert t2.encode("low!")[-1] == 0                # unseen character is still <unk>535455def test_loaded_model_gives_identical_logits(saved):56    prefix, params, cfg, tok = saved57    p2, c2, t2 = load_checkpoint(prefix)58    ids = np.array(tok.encode("low newest tower"))59    a, _ = forward(params, ids, cfg)60    b, _ = forward(p2, ids, c2)61    assert np.array_equal(a, b)626364def test_non_ascii_characters_survive(tmp_path):65    tok = build_tokenizer("café café über\n", max_merges=3)66    cfg = ModelConfig(tok.vocab_size, 8, 4, 2, 8, 1)67    prefix = str(tmp_path / "m")68    save_checkpoint(prefix, init_params(cfg, np.random.default_rng(0)), cfg, tok)69    t2 = load_checkpoint(prefix)[2]70    assert t2.decode(t2.encode("café über")) == "café über"717273def test_tok_file_is_plain_json(saved):74    d = json.load(open(saved[0] + ".tok.json", encoding="utf-8"))75    assert set(d) == {"merges", "itos"}76    assert d["itos"][0] == "<unk>"