rzen
8915d274f5
- src/sample.py: load checkpoint + tokenizer, generate token-by-token with temperature/top-k (or greedy), stop at <|endoftext|>, crop to context window. --suite samples a fixed prompt suite into reports/samples-<config>.md. - tests/test_sample.py: length, context-window safety, early stop, greedy determinism, str output (5 tests, tiny untrained model)
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
"""Sampling tests: generation shape, early stop, greedy determinism.
|
|
|
|
Uses a tiny untrained model (output is gibberish — we only check mechanics).
|
|
Run: `.venv/bin/python tests/test_sample.py`
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT / "src"))
|
|
|
|
from model import GPT # noqa: E402
|
|
from sample import generate, sample_text # noqa: E402
|
|
from tokenizer import BPETokenizer # noqa: E402
|
|
|
|
EOT = "<|endoftext|>"
|
|
|
|
|
|
def _tiny_model():
|
|
cfg = SimpleNamespace(vocab_size=300, context_length=32, n_layers=2,
|
|
n_heads=2, d_model=32, d_ff=64)
|
|
torch.manual_seed(0)
|
|
return GPT(cfg).eval()
|
|
|
|
|
|
def _tiny_tokenizer():
|
|
corpus = ("The cat sat. Tom ran. Lily played happily.\n" + EOT + "\n") * 30
|
|
return BPETokenizer.train(corpus, vocab_size=300, special_tokens=[EOT])
|
|
|
|
|
|
def test_generate_appends_expected_length():
|
|
model = _tiny_model()
|
|
idx = torch.zeros((1, 3), dtype=torch.long)
|
|
out = generate(model, idx, max_new_tokens=10, temperature=0.8, top_k=50)
|
|
# Without an early stop, output grows by exactly max_new_tokens.
|
|
assert out.shape[1] == 3 + 10
|
|
|
|
|
|
def test_generate_respects_context_window():
|
|
model = _tiny_model() # context_length = 32
|
|
idx = torch.zeros((1, 40), dtype=torch.long) # longer than the window
|
|
out = generate(model, idx, max_new_tokens=5, temperature=0.8, top_k=50)
|
|
assert out.shape[1] == 40 + 5 # must not error on over-long context
|
|
|
|
|
|
def test_generate_stops_at_stop_id():
|
|
model = _tiny_model()
|
|
idx = torch.zeros((1, 2), dtype=torch.long)
|
|
# Greedy with stop_id equal to whatever greedy picks first -> stops immediately.
|
|
with torch.no_grad():
|
|
first = model(idx)[0][:, -1, :].argmax().item()
|
|
out = generate(model, idx, max_new_tokens=20, temperature=0.0, stop_id=first)
|
|
assert out.shape[1] == 3 # original 2 + the single stop token
|
|
|
|
|
|
def test_greedy_is_deterministic():
|
|
model = _tiny_model()
|
|
idx = torch.zeros((1, 4), dtype=torch.long)
|
|
a = generate(model, idx, max_new_tokens=8, temperature=0.0)
|
|
b = generate(model, idx, max_new_tokens=8, temperature=0.0)
|
|
assert torch.equal(a, b)
|
|
|
|
|
|
def test_sample_text_returns_str_without_eot():
|
|
model = _tiny_model()
|
|
tok = _tiny_tokenizer()
|
|
text = sample_text(model, tok, "The cat", "cpu",
|
|
max_new_tokens=15, temperature=0.8, top_k=50)
|
|
assert isinstance(text, str)
|
|
assert EOT not in text # the special token is stripped from output
|
|
|
|
|
|
def main() -> None:
|
|
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
|
for t in tests:
|
|
t()
|
|
print(f"ok {t.__name__}")
|
|
print(f"\n{len(tests)} passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|