| 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')) |
|
|
