tests/test_tokenizer.py

"""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!")