Spaces:
Sleeping
Sleeping
File size: 5,562 Bytes
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 | import numpy as np
import pytest
import torch
from plastic.data.mqar import mqar_accuracy, mqar_batch
from plastic.data.physics import PhysicsEnv, physics_batch
from plastic.data.text import TokenWindows, encode_documents_to_bin, split_wikitext_documents
from plastic.tokenizer.bpe import Tokenizer
DOCS = ["alpha beta gamma delta " * 30, "one two three four five " * 30, "the end of it " * 30]
def _tok() -> Tokenizer:
return Tokenizer.train(DOCS, vocab_size=300)
def test_encode_documents_writes_uint16_with_eos(tmp_path):
tok = _tok()
path = str(tmp_path / "train.bin")
n = encode_documents_to_bin(tok, DOCS, path)
arr = np.fromfile(path, dtype="<u2")
assert len(arr) == n
assert int((arr == tok.eos_id).sum()) == len(DOCS)
assert arr[-1] == tok.eos_id
def test_encode_documents_honors_max_tokens(tmp_path):
tok = _tok()
path = str(tmp_path / "small.bin")
n = encode_documents_to_bin(tok, DOCS, path, max_tokens=50)
assert n == 50
assert len(np.fromfile(path, dtype="<u2")) == n
def test_token_windows_sample_and_sequential(tmp_path):
tok = _tok()
path = str(tmp_path / "train.bin")
n = encode_documents_to_bin(tok, DOCS, path)
win = TokenWindows(path, seq_len=16)
g = torch.Generator().manual_seed(0)
x = win.sample(4, g)
assert x.shape == (4, 17) and x.dtype == torch.long
assert int(x.max()) < tok.vocab_size
total = 0
last = None
for batch in win.sequential(batch=2):
assert batch.shape[1] == 17
total += batch.shape[0] * 16
last = batch
assert total <= n and total >= n - 2 * 16 - 1
assert last is not None
def test_split_wikitext_documents():
lines = [
" \n",
" = Article One = \n",
" Some text about one . \n",
" = = Section = = \n",
" More text . \n",
" = Article Two = \n",
" Text about two . \n",
]
docs = list(split_wikitext_documents(lines))
assert len(docs) == 2
assert docs[0].startswith("= Article One =") and "Section" in docs[0]
assert docs[1].startswith("= Article Two =")
def test_mqar_batch_structure():
g = torch.Generator().manual_seed(0)
toks, mask = mqar_batch(3, n_pairs=4, seq_len=32, vocab_size=64, key_range=(3, 30), value_range=(30, 64), rng=g)
assert toks.shape == (3, 32) and mask.shape == (3, 32) and mask.dtype == torch.bool
for b in range(3):
keys = toks[b, 1:9:2].tolist()
vals = toks[b, 2:9:2].tolist()
assert len(set(keys)) == 4 and all(3 <= k < 30 for k in keys) and all(30 <= v < 64 for v in vals)
for t in torch.nonzero(mask[b]).flatten().tolist():
k, v = int(toks[b, t]), int(toks[b, t + 1])
assert vals[keys.index(k)] == v
assert int(mask.sum()) == 3 * 4
def test_mqar_oracle_accuracy():
g = torch.Generator().manual_seed(1)
toks, mask = mqar_batch(2, n_pairs=3, seq_len=20, vocab_size=40, key_range=(3, 20), value_range=(20, 40), rng=g)
logits = torch.full((2, 20, 40), -10.0)
logits.scatter_(2, toks[:, 1:].unsqueeze(-1).expand(-1, -1, 1).clone().roll(0, 1)[:, :, :1].new_zeros(2, 19, 1) + toks[:, 1:].unsqueeze(-1), 10.0) if False else None
# oracle: at every position predict the next token
for b in range(2):
for t in range(19):
logits[b, t, int(toks[b, t + 1])] = 10.0
assert mqar_accuracy(logits, toks, mask) == 1.0
def test_physics_env_matches_closed_form():
env = PhysicsEnv(mu=0.1)
obs = env.reset()
assert obs.shape == (4,)
vel = torch.zeros(2)
pos = torch.zeros(2)
for _ in range(5):
a = torch.randn(2)
obs = env.step(a)
vel = (1 - 0.1) * vel + a
pos = pos + vel
assert torch.allclose(obs, torch.cat([pos, vel]), atol=1e-6)
def test_physics_batch_shapes_resets_and_targets():
g = torch.Generator().manual_seed(0)
batch = physics_batch(3, seq_len=64, episodes_per_seq=4, mu_range=(0.02, 0.25), nonlinear=False, action_std=0.5, rng=g)
assert batch.inputs.shape == (3, 64, 7) and batch.target_delta.shape == (3, 64, 4) and batch.mu.shape == (3, 4)
assert int(batch.inputs[0, :, 6].sum()) == 4
assert batch.inputs[0, 0, 6] == 1 and batch.inputs[0, 16, 6] == 1
obs = batch.inputs[:, :, :4]
inside = batch.inputs[:, 1:, 6] == 0
assert torch.allclose((obs[:, 1:] - obs[:, :-1])[inside], batch.target_delta[:, :-1][inside], atol=1e-5)
assert (obs[:, 16] == 0).all()
@pytest.mark.parametrize("nonlinear", [False, True])
@pytest.mark.parametrize("episodes", [1, 2, 4])
def test_physics_targets_match_env_at_every_step_including_episode_ends(nonlinear, episodes):
"""Every target, including the last step of each episode and of the sequence, is the
transition the action actually produced; resets never leak into targets."""
g = torch.Generator().manual_seed(1)
seq_len = 24
batch = physics_batch(2, seq_len=seq_len, episodes_per_seq=episodes, mu_range=(0.05, 0.3), nonlinear=nonlinear, action_std=0.5, rng=g)
ep_len = seq_len // episodes
for b in range(2):
for e in range(episodes):
env = PhysicsEnv(float(batch.mu[b, e]), nonlinear=nonlinear)
obs = env.reset()
for i in range(ep_len):
t = e * ep_len + i
assert torch.allclose(batch.inputs[b, t, :4], obs, atol=1e-5), (b, e, i)
nxt = env.step(batch.inputs[b, t, 4:6])
assert torch.allclose(batch.target_delta[b, t], nxt - obs, atol=1e-5), (b, e, i)
obs = nxt
|