mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-09 14:03:42 +08:00
YassineYousfi
185 lines
11 KiB
Python
185 lines
11 KiB
Python
import unittest, base64, functools, re, time, unicodedata, json, pathlib, tempfile
|
|
from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
|
|
from tinygrad.helpers import fetch
|
|
|
|
class TestLLMTokenizer(unittest.TestCase):
|
|
@functools.cached_property
|
|
def llama_tok(self):
|
|
# from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
|
|
model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
|
|
with open(model_file, "rt") as fd:
|
|
str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line]
|
|
|
|
# https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
|
|
bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves
|
|
_byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
|
|
_byte_encoder = {v:k for k,v in _byte_decoder.items()}
|
|
normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab}
|
|
|
|
special_tokens = [
|
|
"<|begin_of_text|>",
|
|
"<|end_of_text|>",
|
|
"<|reserved_special_token_0|>",
|
|
"<|reserved_special_token_1|>",
|
|
"<|reserved_special_token_2|>",
|
|
"<|reserved_special_token_3|>",
|
|
"<|start_header_id|>",
|
|
"<|end_header_id|>",
|
|
"<|reserved_special_token_4|>",
|
|
"<|eot_id|>",
|
|
] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
|
|
return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)})
|
|
|
|
def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
|
|
self.assertEqual(tok.encode(text), expected_tokens)
|
|
self.assertEqual(tok.decode(expected_tokens), text)
|
|
|
|
# NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
|
|
def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])
|
|
|
|
def test_llama_basic(self): self._test_coding(self.llama_tok, "hello world", [ 15339, 1917 ])
|
|
def test_llama_control_char(self): self._test_coding(self.llama_tok, " \x850", [ 220, 116360, 15 ])
|
|
def test_llama_bytes(self): self._test_coding(self.llama_tok, " \xec\x8b\xa4\xed", [ 1717, 105, 116174, 82638, 2483 ])
|
|
def test_llama_special1(self): self._test_coding(self.llama_tok, "hello <|end_of_text|>", [ 15339, 220, 128001 ])
|
|
def test_llama_special2(self): self._test_coding(self.llama_tok, "<|start_header_id|>user<|end_header_id|>\n\n", [ 128006, 882, 128007, 271 ])
|
|
def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
|
|
def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n \n", [ 31213, 14211 ])
|
|
|
|
def test_split_regex_matches_naive_listing(self):
|
|
# the compacted codepoint ranges must match the same text as listing every codepoint
|
|
def naive(pre): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
|
|
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + naive("Z"), naive("N"), naive("L")
|
|
naive_re = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" +
|
|
f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
|
|
sample = "hello world 한국어 中文 текст ١٢٣ 123 😊\n \ttoday\n'équivalent ²³№ "
|
|
self.assertEqual(SimpleTokenizer({}, {})._split_to_word.findall(sample), naive_re.findall(sample))
|
|
|
|
def test_split_regex_speed(self):
|
|
# the naive listing compiles a 429KB pattern that takes 10+s to match a 225KB prompt; ranges keep it small and fast
|
|
tok = SimpleTokenizer({}, {})
|
|
self.assertLess(len(tok._split_to_word.pattern), 100_000)
|
|
text = "The quick brown fox jumps over the lazy dog. " * 5000
|
|
tok._split_to_word.findall(text) # warmup
|
|
tms = []
|
|
for _ in range(5):
|
|
st = time.perf_counter()
|
|
words = tok._split_to_word.findall(text)
|
|
tms.append(time.perf_counter() - st)
|
|
self.assertLess(min(tms), 4) # best-of-5 is robust to CI scheduling pauses; new code takes ~60ms
|
|
self.assertEqual(len(words), 50001)
|
|
|
|
def test_llama_continued_conversation(self):
|
|
self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
|
|
self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])
|
|
self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
|
|
|
|
def test_long_cached_prompt_matches_fresh_tokenization(self):
|
|
prefix = "system tools\n" * 700 + "<|eot_id|>"
|
|
first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
|
|
expected = self.llama_tok.encode(changed)
|
|
self.llama_tok.encode(first)
|
|
self.assertEqual(self.llama_tok.encode(changed), expected)
|
|
|
|
def test_tekken_from_gguf_kv(self):
|
|
kv = {
|
|
"tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
|
|
"tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
|
|
"tokenizer.ggml.pre": "tekken",
|
|
"tokenizer.ggml.eos_token_id": 2,
|
|
}
|
|
tok = SimpleTokenizer.from_gguf_kv(kv)
|
|
template = FallbackTemplate(tok)
|
|
self.assertEqual(template.role("user"), "[INST]")
|
|
self.assertEqual(tok.encode("hello"), [5])
|
|
self.assertEqual(template.end_turn(), "[/INST]")
|
|
self.assertEqual(template.role("assistant"), "")
|
|
|
|
def test_gguf_unused_and_eom(self):
|
|
kv = { "tokenizer.ggml.tokens": ["<unk>", "</s>", "[PAD7]", "[", "P", "A", "D", "7", "]"],
|
|
"tokenizer.ggml.token_type": [2, 3, 5, 1, 1, 1, 1, 1, 1],
|
|
"tokenizer.ggml.pre": "llama3", "tokenizer.ggml.eos_token_id": 1, "tokenizer.ggml.eom_token_id": 0 }
|
|
tok = SimpleTokenizer.from_gguf_kv(kv)
|
|
self.assertNotIn(2, tok.encode("[PAD7]")) # UNUSED(5) tokens are not matched as special
|
|
self.assertEqual(tok.decode([2]), "[PAD7]")
|
|
self.assertTrue(tok.is_end(0)) # eom ends generation
|
|
self.assertFalse(tok.is_end(2))
|
|
|
|
def test_hf_tokens(self):
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
path = pathlib.Path(folder)
|
|
vocab = {"h": 0, "i": 1, "hi": 2, "Ġ": 3, "!": 4, "[gMASK]": 5, "hello world": 9, "1": 10, "2": 11, "3": 12, "123": 13}
|
|
added = [dict(content=t, id=i, special=i != 9) for t, i in (("[gMASK]", 5), ("<eos>", 6), ("<eot>", 7), ("<eom>", 8), ("hello world", 9))]
|
|
(path / "tokenizer.json").write_text(json.dumps(dict(model=dict(vocab=vocab), added_tokens=added)))
|
|
(path / "config.json").write_text(json.dumps(dict(eos_token_id=[6, 7, 8])))
|
|
for preset in ("glm4", "llama3", "tekken"):
|
|
with self.subTest(preset=preset):
|
|
tok = SimpleTokenizer.from_hf(str(path), preset=preset)
|
|
self.assertEqual(tok.preset, preset)
|
|
self._test_coding(tok, "hi hi!", [2, 3, 2, 4])
|
|
self._test_coding(tok, "h!i", [0, 4, 1])
|
|
self._test_coding(tok, "hi [gMASK]hello world!<eos>", [2, 3, 5, 9, 4, 6])
|
|
self._test_coding(tok, "123", [10, 11, 12] if preset == "tekken" else [13])
|
|
self.assertEqual((tok.eos_id, tok.eot_id, tok.eom_id), (6, 7, 8))
|
|
self.assertEqual([tok.is_end(i) for i in range(5, 10)], [False, True, True, True, False])
|
|
|
|
def test_hf_metadata(self):
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
path = pathlib.Path(folder)
|
|
(path / "tokenizer.json").write_text(json.dumps(dict(model=dict(vocab={"<s>": 0}))))
|
|
(path / "config.json").write_text(json.dumps(dict(eos_token_id=1)))
|
|
tok = SimpleTokenizer.from_hf(path, preset="llama3")
|
|
self.assertIsNone(tok.bos_id)
|
|
self.assertIsNone(tok.chat_template)
|
|
self.assertEqual((tok.eos_id, tok.eot_id, tok.eom_id), (1, None, None))
|
|
self.assertTrue(tok.is_end(1))
|
|
(path / "tokenizer_config.json").write_text(json.dumps(dict(bos_token="<s>", chat_template="config template")))
|
|
tok = SimpleTokenizer.from_hf(path, preset="llama3")
|
|
self.assertEqual(tok.bos_id, 0)
|
|
self.assertEqual(tok.chat_template, "config template")
|
|
(path / "chat_template.jinja").write_text("{{ messages[0].content }}")
|
|
self.assertEqual(SimpleTokenizer.from_hf(path, preset="llama3").chat_template, "{{ messages[0].content }}")
|
|
|
|
def test_tekken_gpt4o_split(self):
|
|
split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o", "kimi-k2")}
|
|
shared = {
|
|
"HelloWorld": ["Hello", "World"],
|
|
" ÜNICODE": [" ÜNICODE"], # Ü: non-ascii upper joins the run
|
|
"é café": ["é", " café"], # first é is e + U+0301 combining acute (NFD)
|
|
"เพื่อน วิ": ["เพื่อน", " วิ"], # thai vowel marks stay in the word
|
|
"a/b\r\n x": ["a", "/b", "\r\n", " x"], # punct tail eats /
|
|
}
|
|
for s, want in shared.items():
|
|
self.assertEqual(split["tekken"](s), want, f"tekken {s!r}")
|
|
self.assertEqual(split["gpt-4o"](s), want, f"gpt-4o {s!r}")
|
|
self.assertEqual(split["kimi-k2"](s), want, f"kimi-k2 {s!r}")
|
|
differ = [
|
|
("12345", list("12345"), ["123", "45"]), # digits: tekken single, o200k groups {1,3}
|
|
("it's I'M don'T", ["it", "'s", " I", "'M", " don", "'T"], ["it's", " I'M", " don'T"]), # contraction: o200k inline suffix
|
|
]
|
|
for s, tk, go in differ:
|
|
self.assertEqual(split["tekken"](s), tk, f"tekken {s!r}")
|
|
self.assertEqual(split["gpt-4o"](s), go, f"gpt-4o {s!r}")
|
|
self.assertEqual(split["kimi-k2"](s), go, f"kimi-k2 {s!r}")
|
|
self.assertEqual(split["kimi-k2"]("a.\n/b"), ["a", ".\n", "/b"])
|
|
|
|
def test_joyai_llm_split(self):
|
|
tok = SimpleTokenizer({}, {}, "joyai-llm")
|
|
tok._encode_word = lambda word: [word.decode()]
|
|
self.assertEqual(tok.encode("é café"), ["é", " café"])
|
|
self.assertEqual(tok.encode(" 12345"), [" ", "123", "45"])
|
|
self.assertEqual(tok.encode("it's I'M don'T"), ["it", "'s", " I", "'M", " don", "'T"])
|
|
self.assertEqual(tok.encode(" 你好hello"), [" ", "你好", "hello"])
|
|
self.assertEqual(tok.encode("a\u200b"), ["a", "\u200b"])
|
|
|
|
def test_stream_decoder(self):
|
|
"""stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
|
|
bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
|
|
be = {b: chr(b) for b in bs} | {b: chr(256+i) for i,b in enumerate(b for b in range(256) if b not in bs)}
|
|
token_bytes = {25677: b'\x20\xf0\x9f\x98', 138: b'\x8a'} # ' ' + 3/4 emoji | 1/4 emoji (qwen3.5)
|
|
tok = SimpleTokenizer({"".join(be[b] for b in v): k for k, v in token_bytes.items()}, {})
|
|
dec = tok.stream_decoder()
|
|
self.assertEqual(dec(25677) + dec(138) + dec(), " 😊")
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|