Spaces:
Sleeping
Sleeping
File size: 6,834 Bytes
595cac0 4b22369 595cac0 | 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | import os
import torch
from plastic.config import ModelConfig
from plastic.data.text import encode_documents_to_bin
from plastic.store import ArtifactStore
from plastic.tokenizer.bpe import Tokenizer
from plastic.train.loop import TrainConfig, train
from plastic.train.schedule import lr_scale
DOCS = ["alpha beta gamma delta epsilon " * 40, "one two three four five six " * 40, "the end of it all " * 40]
def test_lr_scale_shape():
assert lr_scale(1, warmup=10, total=100) == 0.1
assert lr_scale(10, warmup=10, total=100) == 1.0
assert abs(lr_scale(100, warmup=10, total=100, min_ratio=0.1) - 0.1) < 1e-9
mid = lr_scale(55, warmup=10, total=100, min_ratio=0.1)
assert 0.1 < mid < 1.0
def _text_corpus(tmp_path) -> tuple[str, Tokenizer]:
tok = Tokenizer.train(DOCS, vocab_size=300)
d = str(tmp_path / "data")
os.makedirs(d, exist_ok=True)
tok.save(os.path.join(d, "tokenizer.json"))
encode_documents_to_bin(tok, DOCS, os.path.join(d, "train.bin"))
encode_documents_to_bin(tok, DOCS[:1], os.path.join(d, "validation.bin"))
return d, tok
def test_text_training_writes_artifacts(tmp_path):
d, tok = _text_corpus(tmp_path)
root = str(tmp_path / "artifacts")
cfg = TrainConfig(
domain="text",
model=ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=16, vocab_size=tok.vocab_size),
artifacts_root=root,
data_dir=d,
steps=6,
batch_size=2,
seq_len=32,
warmup_steps=2,
eval_every=3,
eval_batches=2,
save_every=3,
log_every=2,
mqar_pairs=(4,),
device="cpu",
)
mid = train(cfg, log=lambda s: None)
store = ArtifactStore(root)
rec = store.load_model_record(mid)
assert rec["status"] == "completed" and rec["steps"] == 6
assert os.path.exists(store.tokenizer_path(mid))
ev = store.read_eval(mid)
assert ev is not None
for key in ("heldout_loss", "heldout_loss_beta0", "memory_value", "mqar_accuracy", "beta_hist", "step"):
assert key in ev, key
assert "4" in ev["mqar_accuracy"]
assert len(ev["beta_hist"]["counts"]) == 20
# bits-per-byte is reported and is the vocab-independent restatement of the held-out loss
import math
assert ev["bytes_per_token"] > 0 and ev["heldout_bpb"] > 0
assert abs(ev["heldout_bpb"] - ev["heldout_loss"] / (math.log(2.0) * ev["bytes_per_token"])) < 1e-9
log = store.read_log(mid)
assert any(r.get("event") == "eval" for r in log) and any("loss" in r and "tok_per_s" in r for r in log)
cfg2, model, info = store.load_checkpoint(mid)
assert info["step"] == 6 and cfg2.vocab_size == tok.vocab_size
def test_text_training_is_deterministic(tmp_path):
d, tok = _text_corpus(tmp_path)
losses = []
for i in range(2):
cfg = TrainConfig(
domain="text",
model=ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=16, vocab_size=tok.vocab_size),
artifacts_root=str(tmp_path / f"a{i}"),
data_dir=d,
steps=3,
batch_size=2,
seq_len=32,
eval_every=0,
save_every=0,
log_every=1,
device="cpu",
mqar_frac=0.0,
)
mid = train(cfg, log=lambda s: None)
losses.append([r["loss"] for r in ArtifactStore(cfg.artifacts_root).read_log(mid) if "loss" in r])
assert losses[0] == losses[1]
def test_physics_training_writes_artifacts(tmp_path):
root = str(tmp_path / "artifacts")
cfg = TrainConfig(
domain="physics",
model=ModelConfig(domain="physics", d_model=32, n_heads=2, n_layers=1, chunk=16),
artifacts_root=root,
steps=4,
batch_size=2,
seq_len=32,
episodes_per_seq=2,
warmup_steps=1,
eval_every=2,
eval_batches=1,
save_every=2,
log_every=1,
device="cpu",
use_muon=False,
)
mid = train(cfg, log=lambda s: None)
store = ArtifactStore(root)
assert store.load_model_record(mid)["status"] == "completed"
ev = store.read_eval(mid)
assert ev is not None and "memory_value" in ev and "beta_hist" in ev
assert mid.startswith("phys_")
def test_domain_mismatch_is_fixed_up(tmp_path):
cfg = TrainConfig(
domain="physics",
model=ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=16),
artifacts_root=str(tmp_path / "a"),
steps=1,
batch_size=1,
seq_len=16,
episodes_per_seq=1,
eval_every=0,
save_every=0,
eval_batches=1,
device="cpu",
use_muon=False,
)
mid = train(cfg, log=lambda s: None)
cfg2, model, _ = ArtifactStore(cfg.artifacts_root).load_checkpoint(mid)
assert cfg2.domain == "physics"
assert isinstance(model, torch.nn.Module)
def test_heldout_loss_is_batch_size_invariant(tmp_path):
from plastic.data.text import TokenWindows
from plastic.model.lm import PlasticLM
from plastic.train.loop import evaluate_text
d, tok = _text_corpus(tmp_path)
cfg_model = ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=16, vocab_size=tok.vocab_size)
torch.manual_seed(0)
lm = PlasticLM(cfg_model)
heldout = TokenWindows(f"{d}/train.bin", seq_len=32)
n_windows = (len(heldout) - 1) // 32
results = []
for bs in (1, 2, 5):
cfg = TrainConfig(domain="text", model=cfg_model, batch_size=bs, seq_len=32, eval_batches=10**6, mqar_pairs=(), device="cpu")
results.append(evaluate_text(lm, cfg, heldout, torch.device("cpu")))
assert all(r["heldout_tokens"] == n_windows * 32 for r in results)
for r in results[1:]:
assert abs(r["heldout_loss"] - results[0]["heldout_loss"]) < 1e-4
assert abs(r["heldout_loss_beta0"] - results[0]["heldout_loss_beta0"]) < 1e-4
def test_adversarial_training_logs_damage(tmp_path):
d, tok = _text_corpus(tmp_path)
root = str(tmp_path / "adv")
# the validation fixture must be long enough for the default canary probes
encode_documents_to_bin(tok, DOCS * 4, os.path.join(d, "validation.bin"))
cfg = TrainConfig(
domain="text",
model=ModelConfig(d_model=32, n_heads=2, n_layers=1, chunk=8, vocab_size=tok.vocab_size),
artifacts_root=root, data_dir=d, steps=4, batch_size=2, seq_len=32, warmup_steps=1, eval_every=0,
save_every=0, eval_batches=1, log_every=1, device="cpu", mqar_frac=0.0,
adversarial=True, adv_every=2, adv_steps=2, adv_suffix_len=8, adv_prefix_len=16,
)
mid = train(cfg, log=lambda s: None)
log = ArtifactStore(root).read_log(mid)
adv = [r for r in log if "adv_damage" in r]
assert len(adv) == 2 and all(r["adv_damage"] == r["adv_damage"] for r in adv)
assert os.path.exists(ArtifactStore(root).canary_path(mid))
|