Spaces:
Runtime error
Runtime error
| from .train import BPETrainer | |
| from .tokenizer import AbyssTokenizer | |
| SAMPLE = ( | |
| "the cat sat on the mat. the cat ran. " | |
| "a dog sat on the log. the dog ran fast! " | |
| ) * 50 | |
| def test_roundtrip(): | |
| trainer = BPETrainer(vocab_size=400, verbose=False) | |
| merges, _ = trainer.train(SAMPLE) | |
| tok = AbyssTokenizer(merges, vocab_size=400) | |
| cases = [ | |
| "the cat sat", | |
| "Hello, world!", | |
| "café résumé naïve", | |
| "emoji 🐈 and 🔥", | |
| "def f(x): return x + 1", | |
| " leading spaces", | |
| "", | |
| ] | |
| for s in cases: | |
| assert tok.decode(tok.encode(s)) == s, f"round-trip failed: {s!r}" | |
| print("roundtrip: OK") | |
| def test_specials(): | |
| trainer = BPETrainer(vocab_size=400) | |
| merges, _ = trainer.train(SAMPLE) | |
| tok = AbyssTokenizer(merges, vocab_size=400) | |
| s = "<|user|>hi<|endofturn|>" | |
| ids = tok.encode(s) | |
| assert tok.special["<|user|>"] in ids | |
| assert tok.special["<|endofturn|>"] in ids | |
| assert tok.decode(ids) == s | |
| print("specials: OK") | |
| def test_compression(): | |
| trainer = BPETrainer(vocab_size=400) | |
| merges, _ = trainer.train(SAMPLE) | |
| tok = AbyssTokenizer(merges, vocab_size=400) | |
| text = "the cat sat on the mat" | |
| n_bytes = len(text.encode("utf-8")) | |
| n_tokens = len(tok.encode(text)) | |
| print(f"compression: {n_bytes} bytes -> {n_tokens} tokens") | |
| assert n_tokens < n_bytes | |
| if __name__ == "__main__": | |
| test_roundtrip() | |
| test_specials() | |
| test_compression() | |
| print("\nall tests passed") |