File size: 7,501 Bytes
b9301f1 | 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 | """Tests for the atom-native input contract and adapter."""
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
import torch
from src.atom_native import AtomNativeModel, FieldAmplitudeController
from src.io.atomizer import Atomizer
class AtomizerTests(unittest.TestCase):
def test_round_trip_is_exact_for_utf8(self) -> None:
text = "ATOM: déjà vu — 42\nfluide."
atomizer = Atomizer(max_span_bytes=16, pack_mode="linguistic")
packets = atomizer.encode(text)
reconstructed = b"".join(packet.payload for packet in packets).decode("utf-8")
self.assertEqual(reconstructed, text)
self.assertGreater(len(packets), 1)
self.assertTrue(all(packet.features.shape == (320,) for packet in packets))
def test_same_surface_span_can_depend_on_context(self) -> None:
# Eager mode yields word-level "world " packets for this invariant.
first = Atomizer(max_span_bytes=16, pack_mode="eager").encode("a world ")
second = Atomizer(max_span_bytes=16, pack_mode="eager").encode("b world ")
first_world = next(packet for packet in first if packet.payload == b"world ")
second_world = next(packet for packet in second if packet.payload == b"world ")
self.assertFalse(torch.allclose(first_world.features, second_world.features))
def test_chunked_stream_is_equivalent_and_lossless(self) -> None:
text = "alpha beta — gamma\n delta"
one_shot = Atomizer(max_span_bytes=8, pack_mode="linguistic").encode(text)
chunked = list(Atomizer(max_span_bytes=8, pack_mode="linguistic").stream(["alpha ", "beta ", "— ", "gamma\n", " delta"]))
self.assertEqual(b"".join(packet.payload for packet in one_shot), text.encode("utf-8"))
self.assertEqual(b"".join(packet.payload for packet in chunked), text.encode("utf-8"))
def test_state_restore_preserves_next_packet_features(self) -> None:
source = Atomizer(max_span_bytes=16, pack_mode="eager")
source.encode("prefix ")
saved = source.state_dict()
expected = source.packet_from_payload(b"next ")
restored = Atomizer(max_span_bytes=16, pack_mode="eager")
restored.load_state_dict(saved)
actual = restored.packet_from_payload(b"next ")
torch.testing.assert_close(expected.features, actual.features)
self.assertEqual(expected.start, actual.start)
self.assertEqual(expected.end, actual.end)
def test_pending_buffer_is_checkpointable(self) -> None:
source = Atomizer(max_span_bytes=16, pack_mode="eager")
source.encode_bytes(b"partial", reset=True, flush=False)
saved = source.state_dict()
expected = source.encode_bytes(b" end ", reset=False, flush=True)
restored = Atomizer(max_span_bytes=16, pack_mode="eager")
restored.load_state_dict(saved)
actual = restored.encode_bytes(b" end ", reset=False, flush=True)
self.assertEqual([packet.payload for packet in expected], [packet.payload for packet in actual])
def test_linguistic_spans_are_longer_and_lossless(self) -> None:
text = "Utilisateur: Bonjour\nAssistant: Salut mon ami, comment vas-tu ?\n"
eager = Atomizer(max_span_bytes=32, pack_mode="eager").encode(text)
ling = Atomizer(max_span_bytes=32, pack_mode="linguistic").encode(text)
self.assertEqual(b"".join(p.payload for p in ling).decode("utf-8"), text)
self.assertGreater(sum(len(p.payload) for p in ling) / len(ling),
sum(len(p.payload) for p in eager) / len(eager))
self.assertLess(len(ling), len(eager))
# Prefer multi-byte spans (not only 1–3 byte scraps)
self.assertGreater(max(len(p.payload) for p in ling), 8)
# Newlines still bound dialogue lines
self.assertTrue(any(p.boundary == "newline" for p in ling))
def test_linguistic_cuts_prefer_whitespace(self) -> None:
# Long alpha run with spaces: cuts should land on spaces, not mid-word.
text = "alpha beta gamma delta epsilon zeta eta theta "
packets = Atomizer(max_span_bytes=16, pack_mode="linguistic").encode(text)
self.assertEqual(b"".join(p.payload for p in packets).decode("utf-8"), text)
for packet in packets[:-1]:
if packet.boundary in {"whitespace", "punctuation"}:
self.assertTrue(
packet.payload.endswith(b" ") or packet.payload[-1:] in b".,!?;:",
msg=repr(packet.payload),
)
class AtomNativeModelTests(unittest.TestCase):
def test_field_controller_caps_only_large_rms(self) -> None:
controller = FieldAmplitudeController(max_rms=2.0)
field, info = controller(torch.ones(8, 8) * 10.0)
self.assertLessEqual(info["field_rms_after"], 2.001)
self.assertLess(info["field_scale"], 1.0)
small, small_info = controller(torch.ones(8, 8) * 0.5)
torch.testing.assert_close(small, torch.ones(8, 8) * 0.5)
self.assertEqual(small_info["field_scale"], 1.0)
def test_dynamics_parameter_projection(self) -> None:
model = AtomNativeModel(
d_model=4,
n_modes=4,
n_atoms_max=32,
max_payload_bytes=12,
field_max_rms=2.0,
energy_decay_bounds=(0.9, 0.999),
atomizer=Atomizer(max_span_bytes=12),
)
with torch.no_grad():
model.core.dynamics.dynamics.energy_decay.fill_(2.0)
model.core.dynamics.dynamics.coupling_scale.fill_(4.0)
model.core.dynamics.dynamics.phase_sync.fill_(3.0)
state = model.stabilize_dynamics_parameters()
self.assertAlmostEqual(state["energy_decay"], 0.999, places=5)
self.assertAlmostEqual(state["coupling_scale"], 2.0, places=5)
self.assertAlmostEqual(state["phase_sync"], 1.0, places=5)
def test_forward_loss_backward_and_checkpoint_reload(self) -> None:
atomizer = Atomizer(max_span_bytes=12)
packets = atomizer.encode("alpha beta gamma delta epsilon.")
self.assertGreaterEqual(len(packets), 3)
model = AtomNativeModel(
d_model=4,
n_modes=4,
n_atoms_max=32,
max_payload_bytes=12,
atomizer=Atomizer(max_span_bytes=12),
)
optimizer = torch.optim.AdamW(model.trainable_parameters, lr=1e-3)
loss, info = model.transition_loss(packets[0], packets[1])
self.assertTrue(torch.isfinite(loss))
self.assertTrue(torch.isfinite(info["surface"]["byte_logits"]).all())
loss.backward()
optimizer.step()
self.assertGreater(len(model.core.atoms), 0)
self.assertTrue(all(torch.isfinite(parameter).all() for parameter in model.parameters()))
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "atom_native.pt"
model.save(path)
reloaded = AtomNativeModel(
d_model=4,
n_modes=4,
n_atoms_max=32,
max_payload_bytes=12,
atomizer=Atomizer(max_span_bytes=12),
)
reloaded.load(path)
self.assertEqual(len(reloaded.core.atoms), len(model.core.atoms))
self.assertEqual(reloaded.core.state.n_modes, 4)
self.assertTrue(torch.isfinite(reloaded.core.state.alpha).all())
if __name__ == "__main__":
unittest.main()
|