| import random |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from grimboost.tokenization.byte_bpe import ByteBPETokenizer, train_byte_bpe |
|
|
| CORPUS = [ |
| "The quick brown fox jumps over the lazy dog. ", |
| "def f(x):\n return x + 1\n", |
| "中文测试：字节级分词必须精确往返。", |
| "Emoji 🙂🚀 and accents café naïve. ", |
| "Special markers <|im_start|>user<|im_end|> are ordinary text here. ", |
| "Tabs\tand multiple spaces and\nnewlines.\n", |
| ] |
|
|
|
|
| def _iter_corpus(): |
| for _ in range(200): |
| for line in CORPUS: |
| yield line |
|
|
|
|
| @pytest.fixture(scope="module") |
| def tokenizer(tmp_path_factory: pytest.TempPathFactory) -> ByteBPETokenizer: |
| path = tmp_path_factory.mktemp("tok") / "bpe.json" |
| return train_byte_bpe( |
| _iter_corpus(), |
| vocab_size=320, |
| save_path=path, |
| min_frequency=1, |
| ) |
|
|
|
|
| def test_roundtrip_all_examples(tokenizer: ByteBPETokenizer): |
| for text in CORPUS: |
| ids = tokenizer.encode(text) |
| assert tokenizer.decode_bytes(ids) == text.encode("utf-8") |
| assert tokenizer.decode(ids) == text |
|
|
|
|
| def test_encode_with_spans(tokenizer: ByteBPETokenizer): |
| for text in CORPUS: |
| data = text.encode("utf-8") |
| ids, spans = tokenizer.encode_with_spans(data) |
| assert spans[0][0] == 0 |
| assert spans[-1][1] == len(data) |
| for (a, b), (c, _) in zip(spans, spans[1:]): |
| assert b == c |
| assert b"".join(tokenizer.token_bytes(i) for i in ids) == data |
|
|
|
|
| def test_random_unicode_roundtrip(tokenizer: ByteBPETokenizer): |
| rng = random.Random(0) |
| alphabet = "abc XYZ\n\t.,!?é中🙂🚀0123456789" |
| for _ in range(50): |
| text = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 80))) |
| ids = tokenizer.encode(text) |
| assert tokenizer.decode_bytes(ids) == text.encode("utf-8") |
|
|
|
|
| def test_save_load(tokenizer: ByteBPETokenizer, tmp_path: Path): |
| path = tmp_path / "saved.json" |
| tokenizer.save(path) |
| loaded = ByteBPETokenizer.load(path) |
| text = "roundtrip café 中文 🙂" |
| assert loaded.decode_bytes(loaded.encode(text)) == text.encode("utf-8") |
| assert loaded.vocab_size == tokenizer.vocab_size |
|
|
|
|
| def test_special_tokens_have_no_bytes(tokenizer: ByteBPETokenizer): |
| for idx in tokenizer.special_ids: |
| assert tokenizer.token_bytes(idx) is None |
|
|
|
|
| def test_encode_bytes_exact_with_invalid_tail(tokenizer: ByteBPETokenizer): |
| samples = [ |
| b"hello \xe4\xb8", |
| b"abc\xff\xfe", |
| b"abc", |
| b"\xf0\x9f\x99", |
| ] |
| for data in samples: |
| ids, spans = tokenizer.encode_bytes_exact(data) |
| recon = b"".join(tokenizer.token_bytes(i) for i in ids) |
| assert recon == data |
| assert spans[0][0] == 0 |
| assert spans[-1][1] == len(data) |
|
|
