File size: 10,120 Bytes
a4b6ee8 | 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 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | """Tests for HV-Word2Vec, Soft Phi, and Multi-Level Query."""
from __future__ import annotations
import pytest
import numpy as np
from palimseste import hv
from palimseste.learner import Encoder
from palimseste.hv_word2vec import HVWord2Vec, Word2VecConfig
from palimseste.soft_phi import SoftPhi, SoftPhiConfig
from palimseste.multiquery import MultiLevelQuery, MultiLevelConfig
from palimseste.memory import Memory
from palimseste.phi import Phi, KernelConfig
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
# ================================================================ HV-Word2Vec
class TestHVWord2Vec:
def _train(self, text, D=3000, epochs=3):
cfg = Word2VecConfig(D=D, n_epochs=epochs, min_word_freq=2)
w2v = HVWord2Vec(config=cfg, rng=np.random.default_rng(42))
w2v.train(text)
return w2v
def test_train_creates_embeddings(self):
text = "the cat sat on the mat. the dog sat on the rug. the cat ran. the dog ran."
w2v = self._train(text)
assert w2v.vocab_size > 0
assert w2v.is_trained
def test_words_in_similar_contexts_are_similar(self):
"""cat and dog both appear with 'sat on the' → should be similar."""
text = ("the cat sat on the mat. the dog sat on the rug. "
"the cat sat on the bed. the dog sat on the floor. "
"the cat ran fast. the dog ran fast. "
"the cat ate food. the dog ate food.")
w2v = self._train(text, D=5000, epochs=5)
sim_cat_dog = w2v.similarity("cat", "dog")
sim_cat_mat = w2v.similarity("cat", "mat")
# cat and dog share contexts → more similar than cat and mat
assert sim_cat_dog > sim_cat_mat, \
f"cat-dog sim {sim_cat_dog:.3f} should be > cat-mat {sim_cat_mat:.3f}"
def test_unknown_word_returns_none(self):
w2v = self._train("hello world hello world")
assert w2v.get_word_hv("xyzunknown") is None
def test_encode_text_semantic(self):
w2v = self._train("hello world hello world test test")
h = w2v.encode_text_semantic("hello world")
assert isinstance(h, hv.HV)
assert h.D == 3000
def test_most_similar(self):
text = ("cat sat on mat. dog sat on mat. cat ran. dog ran. "
"bird flew. bird sat on mat.")
w2v = self._train(text, D=5000, epochs=5)
similar = w2v.most_similar("cat", top_k=5)
assert len(similar) > 0
# dog should be in the top results (shared context with cat)
words = [w for w, _ in similar]
assert "dog" in words
def test_save_load(self, tmp_path):
w2v = self._train("test test hello hello world world")
w2v.save(str(tmp_path / "w2v"))
w2v2 = HVWord2Vec(config=Word2VecConfig(D=3000), rng=np.random.default_rng(0))
w2v2.load(str(tmp_path / "w2v"))
assert w2v2.vocab_size == w2v.vocab_size
assert w2v2.get_word_hv("test") == w2v.get_word_hv("test")
def test_config_validation(self):
with pytest.raises(ValueError):
Word2VecConfig(D=0)
with pytest.raises(ValueError):
Word2VecConfig(window_radius=0)
# ================================================================ Soft Phi
class TestSoftPhi:
def _setup(self, D=3000):
mem = Memory(D=D, rng=np.random.default_rng(0))
rng = np.random.default_rng(1)
# store 3 traces with known values
for _ in range(3):
a = hv.random_hv(D, rng=rng)
v = hv.random_hv(D, rng=rng)
mem.write(a, v)
return mem, rng
def test_soft_phi_returns_hv(self):
mem, rng = self._setup()
phi = SoftPhi(config=SoftPhiConfig(temperature=10.0))
q = hv.random_hv(3000, rng=rng)
result = phi(mem, q)
# might be None if no candidates, or an HV
if result is not None:
assert isinstance(result, hv.HV)
assert result.D == 3000
def test_soft_phi_blends_results(self):
"""Soft Phi should produce a blend, not just the nearest match."""
D = 5000
mem = Memory(D=D, rng=np.random.default_rng(0))
rng = np.random.default_rng(1)
# create a query that's roughly between two stored addresses
addr1 = hv.random_hv(D, rng=rng)
addr2 = hv.random_hv(D, rng=rng)
val1 = hv.random_hv(D, rng=rng)
val2 = hv.random_hv(D, rng=rng)
mem.write(addr1, val1)
mem.write(addr2, val2)
# query = midpoint (bundle of both addresses)
query = hv.bundle([addr1, addr2], rng=rng, deterministic=True)
phi = SoftPhi(config=SoftPhiConfig(temperature=5.0))
result = phi(mem, query)
if result is not None:
# result should be somewhat similar to both val1 and val2
sim1 = hv.similarity(result, val1)
sim2 = hv.similarity(result, val2)
# at least one should be positive (the blend favors the closer one)
assert sim1 > 0 or sim2 > 0, \
f"blend should resemble at least one value: sim1={sim1:.3f} sim2={sim2:.3f}"
def test_soft_phi_empty_memory(self):
mem = Memory(D=2000, rng=np.random.default_rng(0))
phi = SoftPhi()
q = hv.random_hv(2000, rng=np.random.default_rng(0))
assert phi(mem, q) is None
def test_soft_phi_retrieve(self):
mem, rng = self._setup()
phi = SoftPhi()
q = hv.random_hv(3000, rng=rng)
ret = phi.retrieve(mem, q)
assert ret is not None # Retrieval object
assert isinstance(ret.matches, list)
# ================================================================ Multi-Level Query
class TestMultiLevelQuery:
def _setup(self, D=3000):
# train word2vec on a small corpus
text = ("the capital of france is paris. "
"the capital of italy is rome. "
"the capital of germany is berlin. "
"france italy germany are countries. "
"paris rome berlin are cities. "
"capital city country france italy germany paris rome berlin. "
"what is the capital of france. "
"what is the capital of italy. "
"what is the capital of germany.")
w2v = HVWord2Vec(config=Word2VecConfig(D=D, n_epochs=5, min_word_freq=2),
rng=np.random.default_rng(42))
w2v.train(text)
# create memory with traces
enc = Encoder(D=D, rng=np.random.default_rng(42))
mem = Memory(D=D, rng=np.random.default_rng(42))
self_hv = hv.random_hv(D, rng=np.random.default_rng(42))
# store fact: "capital of france" -> "paris"
# at all three levels
for question, answer in [("capital of france", "paris"),
("capital of italy", "rome"),
("capital of germany", "berlin")]:
# surface level
q_hvs = [enc.encode_str(w) for w in question.split()]
q_bound = []
for i, h in enumerate(q_hvs):
q_bound.append(hv.bind(h, enc._role(i)))
q_surface = hv.bundle(q_bound, rng=np.random.default_rng(0), deterministic=True)
a_hv = enc.encode_str(f"__word_{answer}")
mem.write(hv.bind(self_hv, q_surface), a_hv)
# semantic level
q_sem = w2v.encode_text_semantic(question)
mem.write(hv.bind(self_hv, q_sem), a_hv)
return mem, enc, w2v, self_hv
def test_multi_level_query_returns_result(self):
mem, enc, w2v, self_hv = self._setup()
mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
self_hv=self_hv)
# create a simple tokenizer-like object
class FakeTok:
def encode(self, text, add_bos=False, add_eos=False):
return text.split()
def token_hv(self, tid):
return enc.encode_str(str(tid))
phi = Phi(config=KernelConfig(radius=500, min_weight=1e-6))
result = mlq.query("capital of france", FakeTok(), phi)
assert result is not None
assert result.source in ("surface", "semantic", "word", "none")
def test_semantic_level_finds_related(self):
"""The semantic level should find 'capital of spain' related to
stored 'capital of france' because the words are semantically similar."""
mem, enc, w2v, self_hv = self._setup()
mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
self_hv=self_hv)
class FakeTok:
def encode(self, text, add_bos=False, add_eos=False):
return text.split()
def token_hv(self, tid):
return enc.encode_str(str(tid))
phi = Phi(config=KernelConfig(radius=500, min_weight=1e-6))
# "capital of spain" was never stored, but "capital" and "of" are
# semantically related to stored "capital of france"
result = mlq.query("capital of spain", FakeTok(), phi)
# semantic level should have non-zero confidence
assert result.semantic_confidence >= 0.0
def test_three_encodings_different(self):
"""The three encodings should produce different HVs."""
mem, enc, w2v, self_hv = self._setup()
mlq = MultiLevelQuery(mem=mem, encoder=enc, word2vec=w2v,
self_hv=self_hv)
class FakeTok:
def encode(self, text, add_bos=False, add_eos=False):
return text.split()
def token_hv(self, tid):
return enc.encode_str(str(tid))
s = mlq.encode_surface("capital of france", FakeTok())
w = mlq.encode_words("capital of france")
sem = mlq.encode_semantic("capital of france")
# all three should be valid HVs
assert isinstance(s, hv.HV)
assert isinstance(w, hv.HV)
assert isinstance(sem, hv.HV)
# they should not be identical
assert s != w or w != sem
|