File size: 2,339 Bytes
948a05a | 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 | """Test tokenizer: BPE training, encoding, decoding, round-trip."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from splitbit_llm.model.tokenizer import BPETokenizer, BOS_ID, EOS_ID, PAD_ID, UNK_ID
def test_tokenizer_basic():
"""Test basic tokenizer training and encoding."""
tok = BPETokenizer(vocab_size=128)
text = "Hello world! This is a test. Hello world again. The quick brown fox."
tok.train(text, verbose=False)
assert tok.actual_vocab_size > 10, f"Vocab too small: {tok.actual_vocab_size}"
assert tok.actual_vocab_size <= 128, f"Vocab too large: {tok.actual_vocab_size}"
print(f" Vocab size: {tok.actual_vocab_size}")
def test_tokenizer_roundtrip():
"""Test encode → decode round-trip."""
tok = BPETokenizer(vocab_size=256)
text = "Hello world! This is a test. Hello world again. The quick brown fox jumps over the lazy dog."
tok.train(text, verbose=False)
ids = tok.encode("Hello world!", add_bos=True, add_eos=True)
decoded = tok.decode(ids)
assert "Hello world!" in decoded, f"Round-trip failed: {decoded}"
print(f" Encoded: {ids}")
print(f" Decoded: {decoded}")
def test_tokenizer_special_tokens():
"""Test special token IDs."""
assert BOS_ID == 1
assert EOS_ID == 2
assert PAD_ID == 0
assert UNK_ID == 3
def test_tokenizer_save_load(tmp_path=None):
"""Test save and load."""
import tempfile
tok = BPETokenizer(vocab_size=128)
tok.train("Hello world test test test. Foo bar baz.", verbose=False)
with tempfile.NamedTemporaryFile(suffix=".json", delete=False, mode="w") as f:
path = f.name
tok.save(path)
tok2 = BPETokenizer.load(path)
assert tok2.actual_vocab_size == tok.actual_vocab_size
assert tok2.merges == tok.merges
# Clean up
os.unlink(path)
print(f" Save/load OK ({tok.actual_vocab_size} tokens)")
if __name__ == "__main__":
print("Running tokenizer tests...")
test_tokenizer_basic()
print(" ✓ test_tokenizer_basic")
test_tokenizer_roundtrip()
print(" ✓ test_tokenizer_roundtrip")
test_tokenizer_special_tokens()
print(" ✓ test_tokenizer_special_tokens")
test_tokenizer_save_load()
print(" ✓ test_tokenizer_save_load")
print("\nAll tokenizer tests passed!")
|