File size: 7,128 Bytes
87c9f03 | 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 | """Tests for the hierarchical (million-token) context window."""
import pytest
import numpy as np
from palimseste.hierarchical_context import (
HierarchicalContext,
HierarchicalContextConfig,
)
from palimseste.hv import HV, similarity, random_hv
from palimseste.learner import Encoder
from palimseste.tokenizer import CharTokenizer
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig, PRESETS
# ----------------------------------------------------------- fixtures
@pytest.fixture
def encoder():
return Encoder(D=2000)
@pytest.fixture
def tokenizer(encoder):
tok = CharTokenizer(encoder=encoder)
tok.build_vocab("abcdefghijklmnopqrstuvwxyz .,!?0123456789")
return tok
@pytest.fixture
def hc(encoder):
return HierarchicalContext(
D=2000, encoder=encoder,
config=HierarchicalContextConfig(chunk_size=10, local_window=8, top_k_chunks=3),
)
# ----------------------------------------------------------- basic functionality
class TestHierarchicalContextBasic:
def test_empty_state(self, hc):
"""Empty context returns a valid HV."""
state = hc.get_state()
assert isinstance(state, HV)
assert state.D == 2000
def test_ingest_increments_tokens(self, hc, tokenizer):
ids = tokenizer.encode("hello world", add_bos=True, add_eos=True)
hc.ingest(ids, tokenizer)
assert hc.n_tokens == len(ids)
assert hc.n_local_tokens <= 8
def test_chunking(self, hc, tokenizer):
"""Tokens are split into chunks of chunk_size."""
ids = tokenizer.encode("aaaaaaaaaa bbbbbbbbbb cccccccccc", add_bos=True, add_eos=True)
hc.ingest(ids, tokenizer)
hc.flush()
# Should have at least 1 chunk
assert hc.n_chunks >= 1
assert hc.n_tokens == len(ids)
def test_stats(self, hc, tokenizer):
ids = tokenizer.encode("test", add_bos=True, add_eos=True)
hc.ingest(ids, tokenizer)
stats = hc.stats()
assert "n_chunks" in stats
assert "n_tokens" in stats
assert "max_context_tokens" in stats
assert stats["max_context_tokens"] == 4096 * 10
# ----------------------------------------------------------- state quality
class TestHierarchicalContextState:
def test_same_query_same_state(self, hc, tokenizer):
"""Same input + same query produces identical state."""
ids = tokenizer.encode("hello world test context", add_bos=True, add_eos=True)
hc.ingest(ids, tokenizer)
s1 = hc.get_state()
s2 = hc.get_state()
assert similarity(s1, s2) == pytest.approx(1.0, abs=0.01)
def test_different_context_different_state(self, encoder, tokenizer):
"""Different contexts produce different states."""
hc1 = HierarchicalContext(D=2000, encoder=encoder)
hc2 = HierarchicalContext(D=2000, encoder=encoder)
ids1 = tokenizer.encode("alpha beta gamma delta", add_bos=True, add_eos=True)
ids2 = tokenizer.encode("one two three four", add_bos=True, add_eos=True)
hc1.ingest(ids1, tokenizer)
hc2.ingest(ids2, tokenizer)
s1 = hc1.get_state()
s2 = hc2.get_state()
# Should NOT be identical
assert similarity(s1, s2) < 0.95
def test_state_is_valid_hv(self, hc, tokenizer):
ids = tokenizer.encode("test", add_bos=True, add_eos=True)
hc.ingest(ids, tokenizer)
state = hc.get_state()
assert isinstance(state, HV)
assert state.D == 2000
assert state.bits is not None
# ----------------------------------------------------------- scale
class TestHierarchicalContextScale:
def test_large_context(self, encoder, tokenizer):
"""100 chunks of 100 tokens = 10K tokens."""
hc = HierarchicalContext(
D=2000, encoder=encoder,
config=HierarchicalContextConfig(chunk_size=100, local_window=64),
)
valid = list(range(tokenizer.vocab_size))
ids = np.random.default_rng(42).integers(0, tokenizer.vocab_size, size=10000).tolist()
hc.ingest(ids, tokenizer)
hc.flush()
assert hc.n_chunks == 100
assert hc.n_tokens == 10000
def test_chunk_pruning(self, encoder, tokenizer):
"""Old chunks are pruned when max_chunks is exceeded."""
hc = HierarchicalContext(
D=2000, encoder=encoder,
config=HierarchicalContextConfig(chunk_size=5, local_window=5, max_chunks=3),
)
ids = list(range(tokenizer.vocab_size)) * 10 # lots of tokens
hc.ingest(ids[:100], tokenizer)
hc.flush()
assert hc.n_chunks <= 3
def test_retrieval_at_scale(self, encoder, tokenizer):
"""get_state works quickly even with many chunks."""
hc = HierarchicalContext(
D=2000, encoder=encoder,
config=HierarchicalContextConfig(chunk_size=50, local_window=32, top_k_chunks=5),
)
rng = np.random.default_rng(0)
ids = rng.integers(0, tokenizer.vocab_size, size=5000).tolist()
hc.ingest(ids, tokenizer)
hc.flush()
assert hc.n_chunks > 50
state = hc.get_state()
assert isinstance(state, HV)
# ----------------------------------------------------------- LM integration
class TestLongContextIntegration:
def test_preset_1b_long(self):
"""The 1b-long preset has context_window=1M."""
cfg = PRESETS["1b-long"]
assert cfg.context_window == 1_000_000
assert cfg.D == 100_000
def test_enable_long_context(self):
"""enable_long_context activates hierarchical context."""
lm = PalimpsesteForCausalLM(
config=PalimpsesteConfig(D=2000, context_window=1_000_000, kernel_radius=100)
)
lm.build_tokenizer("hello world test context")
assert lm.long_context_stats() is None
lm.enable_long_context(chunk_size=10, local_window=8)
assert lm.long_context_stats() is not None
def test_state_hv_uses_hier(self):
"""_state_hv uses hierarchical context when enabled."""
lm = PalimpsesteForCausalLM(
config=PalimpsesteConfig(D=2000, context_window=1_000_000, kernel_radius=100)
)
lm.build_tokenizer("hello world test context")
lm.enable_long_context(chunk_size=10, local_window=8)
ids = lm.tokenizer.encode("hello world", add_bos=True, add_eos=True)
hv = lm._state_hv(ids)
assert isinstance(hv, HV)
stats = lm.long_context_stats()
assert stats["n_tokens"] > 0
def test_respond_with_long_context(self):
"""respond() works with long context enabled."""
lm = PalimpsesteForCausalLM(
config=PalimpsesteConfig(D=2000, context_window=1_000_000, kernel_radius=100)
)
lm.build_tokenizer("hello world test context system memory hypervector")
lm.enable_long_context(chunk_size=10, local_window=8, top_k_chunks=3)
resp = lm.respond("hello", max_new_tokens=5)
assert isinstance(resp, str)
stats = lm.long_context_stats()
assert stats["n_tokens"] > 0
|