tests/test_tokenizer.py50 linesGitHub: tests/test_tokenizer.py1import pytest23from tokenizer import (4 UNK, count_words, merge_tokens, train_bpe_on_text, build_tokenizer,5)67CORPUS = "low low low lower newest newest"8910def test_count_words_splits_words_spaces_punctuation():11 assert count_words("a a b!") == {"a": 2, " ": 2, "b": 1, "!": 1}121314def test_merge_tokens_simple():15 assert merge_tokens(["a", "b", "a", "b"], ("a", "b"), "ab") == ["ab", "ab"]161718def test_merge_tokens_overlap_is_left_to_right():19 # a a a -> (aa) a, not a (aa)20 assert merge_tokens(["a", "a", "a"], ("a", "a"), "aa") == ["aa", "a"]212223def test_first_three_merges_match_hand_trace():24 merges = train_bpe_on_text(CORPUS, max_merges=3)25 got = [(m["pair"], m["frequency"]) for m in merges]26 assert got == [(("l", "o"), 4), (("lo", "w"), 4), (("n", "e"), 2)]272829@pytest.fixture30def tok():31 return build_tokenizer(CORPUS, max_merges=5)323334def test_unk_is_id_zero(tok):35 assert tok.itos[0] == UNK363738def test_no_characters_dropped(tok):39 ids = tok.encode("lot") # 'lo' + 't'40 assert len(ids) == 241 assert 0 not in ids424344def test_unseen_character_becomes_unk(tok):45 assert tok.encode("low!")[-1] == 0464748@pytest.mark.parametrize("s", ["low newest", "lower lower", "lot", " low "])49def test_round_trip(tok, s):50 assert tok.decode(tok.encode(s)) == s