File size: 5,926 Bytes
8f2c86e | 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 | """Tests for conversational Q/A training and respond() (``palimseste.lm``).
Verifies:
- train_on_qa_pairs writes traces to M (O(1) per token)
- respond() returns a non-empty string after training
- respond() recovers the exact trained answer for a seen question (greedy)
- respond() handles unseen questions gracefully (no crash)
- the episode format (BOS q EOS BOS a EOS) is consistent
- generate_from_ids respects stop_on_eos
"""
from __future__ import annotations
import numpy as np
import pytest
from palimseste.lm import PalimpsesteForCausalLM, PalimpsesteConfig
from palimseste.tokenizer import BOS, EOS
def _lm(D=5000, radius=30, ctx_window=48):
cfg = PalimpsesteConfig(D=D, context_window=ctx_window, kernel_radius=radius,
temperature=0.3)
return PalimpsesteForCausalLM(config=cfg)
# ----------------------------------------------------------------- train
def test_train_on_qa_pairs_grows_memory():
lm = _lm()
pairs = [("hello", "hi there"), ("bye", "goodbye")]
full = "".join(q + a for q, a in pairs)
lm.build_tokenizer(full)
n = lm.train_on_qa_pairs(pairs)
assert n > 0
assert len(lm.mem) == n
def test_train_on_qa_pairs_requires_tokenizer():
lm = _lm()
with pytest.raises(AssertionError):
lm.train_on_qa_pairs([("q", "a")])
# ----------------------------------------------------------------- respond
def test_respond_returns_exact_answer_for_seen_question():
lm = _lm(radius=0)
pairs = [
("bonjour", "salut je suis palimpseste"),
("qui es-tu", "je suis palimpseste"),
("comment tu apprends", "j'apprends en ecrivant"),
]
full = "".join(q + a for q, a in pairs)
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
# greedy: should recover the exact trained answer
resp = lm.respond("bonjour", temperature=0.0, seed=0)
assert resp == "salut je suis palimpseste"
def test_respond_recovers_multiple_answers():
lm = _lm(radius=0)
pairs = [
("hello", "hi i am palimpseste"),
("who are you", "i am palimpseste a hypervectorial cortex"),
("how do you learn", "i learn by writing to memory"),
("do you use a gpu", "no i do not use a gpu"),
]
full = "".join(q + a for q, a in pairs)
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
for q, expected in pairs:
resp = lm.respond(q, temperature=0.0, seed=0)
assert resp == expected, f"Q={q!r}: got {resp!r}, expected {expected!r}"
def test_respond_unseen_question_no_crash():
lm = _lm()
pairs = [("hello", "hi there")]
full = "hellohi there"
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
# unseen question should not crash; may return empty or partial
resp = lm.respond("what is the weather", temperature=0.5, seed=0)
assert isinstance(resp, str)
def test_respond_empty_input_no_crash():
lm = _lm()
lm.build_tokenizer("abc")
lm.train_on_qa_pairs([("a", "b")])
resp = lm.respond("", temperature=0.0, seed=0)
assert isinstance(resp, str)
def test_respond_respects_max_tokens():
lm = _lm(radius=0)
# long answer
pairs = [("q", "a" * 100)]
full = "q" + "a" * 100
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
resp = lm.respond("q", max_new_tokens=5, temperature=0.0, seed=0)
assert len(resp) <= 5
# ----------------------------------------------------------------- generate_from_ids
def test_generate_from_ids_stop_on_eos():
lm = _lm(radius=0)
text = "hello world"
lm.build_tokenizer(text)
lm.train_on_text(text)
ids = lm.tokenizer.encode("hello", add_bos=True)
out = lm.generate_from_ids(ids, max_new_tokens=50, temperature=0.0, seed=0,
stop_on_eos=True)
assert isinstance(out, str)
def test_generate_from_ids_cold_start():
lm = _lm()
lm.build_tokenizer("abc")
# untrained context -> cold start -> empty string
out = lm.generate_from_ids([BOS], max_new_tokens=10, temperature=0.0, seed=0)
assert isinstance(out, str)
# ----------------------------------------------------------------- episode format
def test_qa_episode_format_consistent():
"""The episode BOS q EOS BOS a EOS must be used consistently in
train_on_qa_pairs and respond()."""
lm = _lm(radius=0)
pairs = [("hi", "hello there")]
full = "hihello there"
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
# respond() must encode the question the same way: BOS q EOS BOS
# then generate. The first generated token should be 'h' (start of "hello there")
q_ids = lm.tokenizer.encode("hi", add_bos=True, add_eos=True)
ctx = q_ids + [BOS]
logits = lm._logits(ctx)
import numpy as np
top = int(np.argmax(logits))
# the first answer char is 'h'
assert lm.tokenizer.id2char[top] == "h"
# ----------------------------------------------------------------- robustness
def test_qa_with_french_accents():
lm = _lm(radius=0)
pairs = [("comment ça va", "ça va bien merci")]
full = "comment ça vaça va bien merci"
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
resp = lm.respond("comment ça va", temperature=0.0, seed=0)
assert "ça va bien" in resp or resp == "ça va bien merci"
def test_qa_multiple_pairs_no_interference():
"""Training on multiple Q/A pairs should not cause one answer to bleed
into another (each episode is independent with its own BOS/EOS boundary)."""
lm = _lm(radius=0)
pairs = [
("aaa", "xxx"),
("bbb", "yyy"),
("ccc", "zzz"),
]
full = "".join(q + a for q, a in pairs)
lm.build_tokenizer(full)
lm.train_on_qa_pairs(pairs)
assert lm.respond("aaa", temperature=0.0, seed=0) == "xxx"
assert lm.respond("bbb", temperature=0.0, seed=0) == "yyy"
assert lm.respond("ccc", temperature=0.0, seed=0) == "zzz"
|