ensemble / tests /test_brain_memory.py
thefinalboss's picture
v0.2: global shared BPE tokenizer + persistent central brain memory + brain save/load. BPE lifts quality ceiling (TinyStories 150KB D=5000: next-token acc 94.5% -> 98.5%, latency 112ms -> 28ms, RAM halved). BrainMemory grows via thinking and survives save/load. 75 tests.
22d1ad7 verified
Raw
History Blame Contribute Delete
4.16 kB
"""Tests for the Brain's central persistent memory."""
from __future__ import annotations
import pytest
from ensemble import Brain, Expert, BrainMemory
from palimseste.hv import random_hv, similarity
class TestBrainMemory:
def test_create_empty(self):
bm = BrainMemory(D=2000)
assert bm.n_concepts == 0
assert len(bm) == 0
def test_remember_and_recall(self):
bm = BrainMemory(D=2000)
addr = random_hv(2000, rng=__import__("numpy").random.default_rng(1))
val = random_hv(2000, rng=__import__("numpy").random.default_rng(2))
bm.remember(addr, val, weight=1.0, tag="test")
assert bm.n_concepts == 1
# recall with the exact address -> should retrieve the value
results = bm.recall(addr, top_k=1)
assert len(results) >= 1
assert results[0].value == val
def test_recall_empty_memory(self):
bm = BrainMemory(D=2000)
q = random_hv(2000)
assert bm.recall(q) == []
def test_dimension_mismatch_raises(self):
bm = BrainMemory(D=2000)
bad = random_hv(3000)
with pytest.raises(ValueError):
bm.remember(bad, bad)
def test_reconstruct(self):
bm = BrainMemory(D=2000)
rng = __import__("numpy").random.default_rng(5)
addr = random_hv(2000, rng=rng)
val = random_hv(2000, rng=rng)
bm.remember(addr, val)
recon = bm.reconstruct(addr)
assert recon is not None
assert similarity(recon, val) > 0.5
def test_save_load_round_trip(self, tmp_path):
bm = BrainMemory(D=2000)
rng = __import__("numpy").random.default_rng(7)
for i in range(5):
bm.remember(random_hv(2000, rng=rng), random_hv(2000, rng=rng))
assert bm.n_concepts == 5
bm.save(tmp_path / "mem.bin")
bm2 = BrainMemory.load(tmp_path / "mem.bin", D=2000)
assert bm2.n_concepts == 5
class TestBrainCentralMemory:
def test_brain_creates_memory_on_first_expert(self):
brain = Brain()
assert brain._brain_memory is None
e = Expert.from_text("hello world. " * 20, D=2000)
brain.add_expert(e)
assert brain._brain_memory is not None
assert brain._D == 2000
def test_brain_remember_recall(self):
brain = Brain()
e = Expert.from_text("hello world. " * 20, D=2000)
brain.add_expert(e)
rng = __import__("numpy").random.default_rng(3)
addr = random_hv(2000, rng=rng)
val = random_hv(2000, rng=rng)
brain.remember(addr, val, tag="manual")
assert brain.n_concepts == 1
results = brain.recall(addr)
assert len(results) >= 1
def test_brain_save_load_persists_concepts(self, tmp_path):
brain = Brain()
e = Expert.from_text("hello world. " * 20, D=2000)
brain.add_expert(e)
rng = __import__("numpy").random.default_rng(9)
for _ in range(3):
brain.remember(random_hv(2000, rng=rng), random_hv(2000, rng=rng))
n_before = brain.n_concepts
assert n_before == 3
brain.save(tmp_path / "brain")
brain2 = Brain.load(tmp_path / "brain")
assert brain2.n_concepts == n_before
def test_think_writes_to_central_memory(self):
"""thinking should write concepts into the central memory when they emerge."""
# use overlapping experts so cross-expert concepts can form
a = Expert.from_qa_pairs(
[("what is x", "x is a thing")] * 4, domain="a", D=3000)
b = Expert.from_qa_pairs(
[("what is x", "x is an object")] * 4, domain="b", D=3000)
brain = Brain()
brain.add_expert(a)
brain.add_expert(b)
n0 = brain.n_concepts
brain.think(n_cycles=4)
# thinking ran without crashing; concepts may or may not form
# (depends on coherence), but n_concepts must be accessible
assert brain.n_concepts >= n0
def test_brain_with_D_param(self):
brain = Brain(D=2500)
assert brain._D == 2500
assert brain._brain_memory is not None
assert brain._brain_memory.D == 2500