File size: 7,159 Bytes
3f69d68 | 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 | import random
from pathlib import Path
import pytest
import torch
from lightnest import coder
def test_input_limits_are_explicit():
with pytest.raises(ValueError, match='character'):
coder.encode_say('返' * 257, None)
with pytest.raises(ValueError, match='unsupported'):
coder.encode_say('\U0001f9ea', None)
with pytest.raises(ValueError, match='limit'):
coder.encode_explain(b'\0' * 96, None)
with pytest.raises(ValueError, match='aligned'):
coder.encode_edit(b'\0', '返回第一个数', None)
def test_batch_fails_instead_of_retrying(monkeypatch):
calls = []
def oversize(rng):
calls.append(1)
return [0] * 257, 1
monkeypatch.setattr(coder, '_sample', oversize)
with pytest.raises(ValueError, match='context'):
coder.make_batch(random.Random(0), 1, torch.device('cpu'))
assert len(calls) == 1
with pytest.raises(ValueError, match='positive'):
coder.make_batch(random.Random(0), 0, torch.device('cpu'))
def test_generation_rejects_context_overflow():
model = coder.Coder(dim=8, layers=1, heads=2)
with pytest.raises(ValueError, match='context'):
coder._generate(model, [coder.BOS] * 250, coder.BYTE_IDS, 96)
with pytest.raises(ValueError, match='length'):
model(torch.zeros((1, 257), dtype=torch.long))
def test_checkpoint_exact_vocab(tmp_path):
path = tmp_path / 'small.pt'
model = coder.Coder(dim=8, layers=1, heads=2)
coder.save_model(model, path)
loaded = coder.load_model(path, torch.device('cpu'))
assert loaded.dim == 8
blob = torch.load(path, weights_only=True)
blob['vocab'][0], blob['vocab'][1] = blob['vocab'][1], blob['vocab'][0]
torch.save(blob, path)
with pytest.raises(ValueError, match='incompatible'):
coder.load_model(path, torch.device('cpu'))
def test_explain_file_does_not_read_whole_file(tmp_path, monkeypatch):
path = tmp_path / 'large.bin'
with path.open('wb') as f:
f.truncate(1024 * 1024)
monkeypatch.setattr(Path, 'read_bytes', lambda *a: pytest.fail('unbounded read'))
with pytest.raises(ValueError):
coder.explain_file(None, path)
def test_training_subset_applies_to_all_modes(monkeypatch):
chosen = tuple(coder.JOBS[:2])
names = {job.name for job in chosen}
codes = {job.code for job in chosen}
observed = set()
def say(sentence, program):
observed.add('say')
assert program in codes
assert sentence in {s for job in chosen for s in job.say_train}
return [0], 0
def edit(prior, sentence, program):
observed.add('edit')
assert prior in codes and program in codes
assert sentence in {s for s, name in coder.EDIT_TRAIN if name in names}
return [0], 0
def explain(program, sentence):
observed.add('explain')
assert program in codes
assert sentence in {job.explain for job in chosen}
return [0], 0
monkeypatch.setattr(coder, 'encode_say', say)
monkeypatch.setattr(coder, 'encode_edit', edit)
monkeypatch.setattr(coder, 'encode_explain', explain)
rng = random.Random(10)
for _ in range(100):
coder._sample(rng, chosen)
assert observed == {'say', 'edit', 'explain'}
def test_semantic_augmentation_changes_bytes_preserves_branches():
from lightnest.phrase import by_name, heldout_probes
from lightnest.rv32 import simulate
order = coder.TEMP_REGISTERS[1:] + coder.TEMP_REGISTERS[:1]
for name in ('add', 'max', 'sum', 'range', 'count_neg'):
job = by_name(name)
transformed = coder.augment_program(job.code, register_order=order, nop_prefix=1)
assert transformed != job.code
assert len(transformed) == len(job.code) + 4
assert all(simulate(transformed, xs) == expected
for xs, expected in heldout_probes(name, seed=88212, trials=12))
with pytest.raises(ValueError, match='permute'):
coder.augment_program(by_name('add').code, register_order=(0, 10, 11, 12))
def test_augmentation_bank_is_train_only_and_context_bounded():
withheld = {'abs', 'mul', 'sum_squares', 'count_even', 'last'}
jobs = tuple(job for job in coder.JOBS if job.name not in withheld)
bank = coder.training_augmentations(jobs, 91241)
assert set(bank) == {job.name for job in jobs}
assert set(bank).isdisjoint(withheld)
for job in jobs:
for variant in bank[job.name]:
assert len(variant) <= coder.MAX_PROGRAM_BYTES
ids, _ = coder.encode_explain(variant, job.explain)
assert len(ids) <= coder.MAX_LEN
for target in jobs:
sentence = max(target.say_train, key=len)
ids, _ = coder.encode_edit(variant, '改成' + sentence, target.code)
assert len(ids) <= coder.MAX_LEN
for _ in range(10):
tokens, _ = coder.make_batch(random.Random(15), 8, torch.device('cpu'), jobs, bank)
assert tokens.shape[1] <= coder.MAX_LEN
def test_augmentation_only_changes_training_inputs(monkeypatch):
jobs = tuple(coder.JOBS[:2])
bank = {job.name: (b'\x13\0\0\0' + job.code,) for job in jobs}
canonical = {job.code for job in jobs}
variants = {value for values in bank.values() for value in values}
modes = set()
def say(sentence, target):
modes.add('say')
assert target in canonical
return [0], 0
def edit(prior, sentence, target):
modes.add('edit')
assert prior in variants and target in canonical
return [0], 0
def explain(prior, sentence):
modes.add('explain')
assert prior in variants
return [0], 0
monkeypatch.setattr(coder, 'encode_say', say)
monkeypatch.setattr(coder, 'encode_edit', edit)
monkeypatch.setattr(coder, 'encode_explain', explain)
rng = random.Random(721)
for _ in range(100):
coder._sample(rng, jobs, bank)
assert modes == {'say', 'edit', 'explain'}
def test_augmentation_rejects_unknown_opcodes_and_oversize():
with pytest.raises(ValueError, match='opcode'):
coder.augment_program(b'\xff' * 4)
with pytest.raises(ValueError, match='limit'):
coder.augment_program(bytes.fromhex('13000000') * 23, nop_prefix=2)
def test_checkpoint_size_budget_checked_before_load(tmp_path, monkeypatch):
path = tmp_path / 'large.pt'
with path.open('wb') as out:
out.truncate(100 * 1024**2)
monkeypatch.setattr(torch, 'load', lambda *a, **k: pytest.fail('must reject before tensor load'))
with pytest.raises(ValueError, match='size budget'):
coder.load_model(path, torch.device('cpu'))
def test_checkpoint_shape_budget_checked_before_construction(tmp_path, monkeypatch):
path = tmp_path / 'small.pt'
coder.save_model(coder.Coder(16, 1, 4), path)
data = torch.load(path, weights_only=True)
data['dim'] = 1000000
torch.save(data, path)
monkeypatch.setattr(coder, 'Coder', lambda *a, **k: pytest.fail('must reject before model allocation'))
with pytest.raises(ValueError, match='parameter budget'):
coder.load_model(path, torch.device('cpu'))
|