rzen
2c5fc29592
- src/tokenizer.py: byte-level BPE with GPT-2-style word pre-tokenization, protected <|endoftext|> special token, save/load as inspectable JSON. Simple recount-per-merge trainer over unique words. - tests/test_tokenizer.py: round-trip, special-token integrity, unicode, lossless pre-tokenization, save/load, determinism (8 tests, no pytest dep) - scripts/train_tokenizer.py: train on a corpus sample, report bytes/token
101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
"""Tokenizer tests. Run directly: `.venv/bin/python tests/test_tokenizer.py`
|
|
|
|
No pytest dependency — plain asserts and a tiny runner, to keep the project's
|
|
dependencies limited to torch + numpy.
|
|
"""
|
|
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
|
|
|
from tokenizer import BPETokenizer, pre_tokenize # noqa: E402
|
|
|
|
EOT = "<|endoftext|>"
|
|
|
|
# A small but varied training corpus.
|
|
CORPUS = (
|
|
"Once upon a time, there was a little cat. The cat liked to play.\n"
|
|
"The cat ran and ran. Then the cat found a red ball!\n" + EOT + "\n"
|
|
"Tom and Lily went to the park. They were very happy.\n"
|
|
) * 50
|
|
|
|
|
|
def _fresh_tokenizer() -> BPETokenizer:
|
|
return BPETokenizer.train(CORPUS, vocab_size=400, special_tokens=[EOT])
|
|
|
|
|
|
def test_pretokenize_is_lossless():
|
|
for text in ["Hello, world!", " spaces\tand\nnewlines ", "a", "", "!!!??"]:
|
|
assert "".join(pre_tokenize(text)) == text
|
|
|
|
|
|
def test_roundtrip_on_training_text():
|
|
tok = _fresh_tokenizer()
|
|
for line in CORPUS.split("\n"):
|
|
assert tok.decode(tok.encode(line)) == line
|
|
|
|
|
|
def test_roundtrip_on_unseen_text():
|
|
tok = _fresh_tokenizer()
|
|
unseen = "A brave dog jumped over the fence; zebras watched quietly."
|
|
assert tok.decode(tok.encode(unseen)) == unseen
|
|
|
|
|
|
def test_special_token_is_single_id():
|
|
tok = _fresh_tokenizer()
|
|
ids = tok.encode(f"Hello {EOT} world")
|
|
eot_id = tok.special_tokens[EOT]
|
|
assert ids.count(eot_id) == 1
|
|
# The special token id must not appear from encoding ordinary text.
|
|
assert eot_id not in tok.encode("Hello world")
|
|
assert tok.decode(tok.encode(f"a{EOT}b")) == f"a{EOT}b"
|
|
|
|
|
|
def test_unicode_roundtrip():
|
|
tok = _fresh_tokenizer()
|
|
# Multibyte characters must survive being split across byte-level tokens.
|
|
text = "café naïve \U0001f600"
|
|
assert tok.decode(tok.encode(text)) == text
|
|
|
|
|
|
def test_vocab_size_and_layout():
|
|
tok = _fresh_tokenizer()
|
|
# vocab_size is a target ceiling: a tiny, repetitive corpus saturates
|
|
# (every word merges to a single token) before reaching it, so we may get
|
|
# fewer merges. The derived vocab dict must always match the reported size.
|
|
assert tok.vocab_size <= 400
|
|
assert len(tok.vocab) == tok.vocab_size
|
|
# First 256 ids are the raw bytes.
|
|
assert tok.vocab[65] == b"A"
|
|
|
|
|
|
def test_save_load_roundtrip():
|
|
tok = _fresh_tokenizer()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
path = Path(d) / "tok.json"
|
|
tok.save(path)
|
|
loaded = BPETokenizer.load(path)
|
|
sample = f"The cat and Tom. {EOT}"
|
|
assert loaded.encode(sample) == tok.encode(sample)
|
|
assert loaded.vocab_size == tok.vocab_size
|
|
|
|
|
|
def test_determinism():
|
|
a = _fresh_tokenizer()
|
|
b = _fresh_tokenizer()
|
|
assert a.merges == b.merges
|
|
|
|
|
|
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()
|