File size: 8,407 Bytes
d9865f3 | 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 | """Tests for the killer generation layer (7 fixes)."""
import pytest
import numpy as np
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.killer import (
NgramPredictor, MultiScaleEncoder, NativeRAG,
IterativeRefiner, TemplateExtractor, massive_ingestion_plan,
NgramPrediction, MultiScaleResult, RAGResult, RefinementResult, Template,
)
def _build_model(D=5000, ctx=128, radius=200):
cfg = PalimpsesteConfig(D=D, context_window=ctx, kernel_radius=radius, temperature=0.0)
lm = PalimpsesteForCausalLM(config=cfg)
pairs = [
("hello", "hi i am palimpseste"),
("who are you", "i am palimpseste a hypervectorial cortex"),
("what is python", "python is a programming language"),
("what is the capital of france", "the capital of france is paris"),
("who won the world cup 2018", "france"),
("what is recursion", "recursion is when a function calls itself"),
]
lm.build_tokenizer("".join(q + a for q, a in pairs))
lm.train_on_qa_pairs(pairs)
return lm, pairs
# ================================================================ BPE
class TestBPETokenizer:
def test_bpe_encode_decode(self):
from palimseste.bpe import BPETokenizer
lm, pairs = _build_model()
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
corpus = " ".join(q + a for q, a in pairs)
bpe.train(corpus)
ids = bpe.encode("hello world", add_bos=True, add_eos=True)
assert len(ids) > 0
text = bpe.decode(ids)
assert "hello" in text
def test_bpe_reduces_tokens(self):
from palimseste.bpe import BPETokenizer
lm, pairs = _build_model()
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=500)
bpe.train(" ".join(q + a for q, a in pairs))
char_ids = lm.tokenizer.encode("hello world", add_bos=True, add_eos=True)
bpe_ids = bpe.encode("hello world", add_bos=True, add_eos=True)
# BPE should produce fewer or equal tokens
assert len(bpe_ids) <= len(char_ids)
def test_bpe_vocab_size(self):
from palimseste.bpe import BPETokenizer
lm, _ = _build_model()
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
bpe.train("hello world test")
assert bpe.vocab_size_actual > 4 # at least special tokens + some bytes
def test_lm_accepts_bpe(self):
"""LM should accept a BPE tokenizer via attach_tokenizer."""
from palimseste.bpe import BPETokenizer
lm, pairs = _build_model()
bpe = BPETokenizer(encoder=lm.encoder, vocab_size=300)
bpe.train(" ".join(q + a for q, a in pairs))
lm.attach_tokenizer(bpe)
# vocab_size should be actual, not target
assert lm.config.vocab_size == bpe.vocab_size_actual
# ================================================================ N-GRAM
class TestNgramPredictor:
def test_learn_ngrams(self):
lm, pairs = _build_model()
ngp = NgramPredictor(lm=lm)
ngp.learn_ngrams([q + a for q, a in pairs])
assert len(ngp._bigram_stats) > 0
assert len(ngp._trigram_stats) > 0
def test_predict_next(self):
lm, pairs = _build_model()
ngp = NgramPredictor(lm=lm)
ngp.learn_ngrams([q + a for q, a in pairs])
q_ids = lm.tokenizer.encode("who are you", add_bos=True, add_eos=True) + [1]
pred = ngp.predict_next(q_ids)
assert isinstance(pred, NgramPrediction)
assert len(pred.token_ids) >= 0
def test_predict_bigram(self):
lm, pairs = _build_model()
ngp = NgramPredictor(lm=lm)
ngp.learn_ngrams([q + a for q, a in pairs])
q_ids = lm.tokenizer.encode("who are you", add_bos=True, add_eos=True) + [1]
pred = ngp.predict_bigram(q_ids)
assert isinstance(pred, NgramPrediction)
def test_cold_start(self):
lm, _ = _build_model()
ngp = NgramPredictor(lm=lm)
pred = ngp.predict_next([0] * 100) # unknown context
assert pred.confidence >= 0.0
# ================================================================ MULTI-SCALE
class TestMultiScaleEncoder:
def test_push_returns_result(self):
lm, _ = _build_model()
mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
result = mse.push(5, lm.tokenizer)
assert isinstance(result, MultiScaleResult)
assert result.state_hv.D == lm.config.D
def test_peek_empty(self):
lm, _ = _build_model()
mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
result = mse.peek(lm.tokenizer)
assert result is None # empty at start
def test_peek_after_push(self):
lm, _ = _build_model()
mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
mse.push(5, lm.tokenizer)
result = mse.peek(lm.tokenizer)
assert result is not None
def test_reset(self):
lm, _ = _build_model()
mse = MultiScaleEncoder(D=lm.config.D, encoder=lm.encoder)
mse.push(5, lm.tokenizer)
mse.reset()
result = mse.peek(lm.tokenizer)
assert result is None # reset clears state
# ================================================================ RAG
class TestNativeRAG:
def test_retrieve_context(self):
lm, _ = _build_model()
rag = NativeRAG(lm=lm, top_k=3)
from palimseste.hv import bind, random_hv
q = random_hv(lm.config.D)
boosted, sims = rag.retrieve_context(q)
assert boosted.D == lm.config.D
def test_generate_with_rag(self):
lm, _ = _build_model()
rag = NativeRAG(lm=lm, top_k=3)
result = rag.generate_with_rag("who are you", max_new_tokens=20)
assert isinstance(result, RAGResult)
assert len(result.text) >= 0
# ================================================================ REFINER
class TestIterativeRefiner:
def test_refine_response(self):
lm, _ = _build_model()
refiner = IterativeRefiner(lm=lm, max_passes=2)
result = refiner.refine_response("who are you", max_new_tokens=20)
assert isinstance(result, RefinementResult)
assert result.n_passes == 2
def test_empty_question(self):
lm, _ = _build_model()
refiner = IterativeRefiner(lm=lm, max_passes=1)
result = refiner.refine_response("xyz123unknown", max_new_tokens=10)
assert isinstance(result, RefinementResult)
# ================================================================ TEMPLATES
class TestTemplateExtractor:
def test_extract(self):
extractor = TemplateExtractor()
templates = extractor.extract([
"python is a language",
"gravity is a force",
"the cat is a mammal",
])
assert len(templates) > 0
assert any("is a" in t.pattern for t in templates)
def test_match_template(self):
extractor = TemplateExtractor()
extractor.extract(["python is a language"])
match = extractor.match_template("java is a language")
assert match is not None
def test_no_match(self):
extractor = TemplateExtractor()
match = extractor.match_template("xyz123 random text")
assert match is None
def test_fill_template(self):
extractor = TemplateExtractor()
result = extractor.fill_template("{0} is a {1}", ["python", "language"])
assert result == "python is a language"
def test_capital_pattern(self):
extractor = TemplateExtractor()
extractor.extract([
"the capital of france is paris",
"the capital of japan is tokyo",
])
match = extractor.match_template("the capital of italy is rome")
assert match is not None
assert "capital" in match[0]
# ================================================================ INGESTION PLAN
class TestMassiveIngestion:
def test_plan_structure(self):
plan = massive_ingestion_plan()
assert "priority_1" in plan
assert "total_estimated" in plan
assert "tokens" in plan["total_estimated"]
def test_plan_has_sources(self):
plan = massive_ingestion_plan()
for key in ["priority_1", "priority_2", "priority_3", "priority_4"]:
assert "source" in plan[key]
assert "estimated_tokens" in plan[key]
|