| """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 |
|
|
|
|
| |
| @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), |
| ) |
|
|
|
|
| |
| 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() |
| |
| 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 |
|
|
|
|
| |
| 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() |
| |
| 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 |
|
|
|
|
| |
| 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 |
| 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) |
|
|
|
|
| |
| 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 |
|
|
