wehub-resource-sync
137 lines
5.4 KiB
Python
137 lines
5.4 KiB
Python
"""
|
|
Test the tokenizer: BPE round-trips, special tokens, and conversation rendering.
|
|
Trains a tiny throwaway tokenizer in-process so the test is hermetic
|
|
(no dependency on ~/.cache/nanochat).
|
|
|
|
python -m pytest tests/test_tokenizer.py -v
|
|
"""
|
|
|
|
import pytest
|
|
from nanochat.tokenizer import RustBPETokenizer, SPECIAL_TOKENS
|
|
|
|
# a small corpus is enough to exercise the BPE machinery
|
|
CORPUS = [
|
|
"The quick brown fox jumps over the lazy dog.",
|
|
"hello world, hello tokenizer, hello hello hello",
|
|
"Numbers like 12345 and unicode like naïve café 你好 🙂 should survive.",
|
|
"def f(x):\n return x + 1\n",
|
|
] * 8
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def tokenizer():
|
|
vocab_size = 256 + len(SPECIAL_TOKENS) + 35 # bytes + specials + a few merges
|
|
return RustBPETokenizer.train_from_iterator(iter(CORPUS), vocab_size)
|
|
|
|
|
|
def test_vocab_size(tokenizer):
|
|
assert tokenizer.get_vocab_size() == 256 + len(SPECIAL_TOKENS) + 35
|
|
|
|
|
|
def test_encode_decode_roundtrip(tokenizer):
|
|
for text in ["hello world", "naïve café 你好 🙂", "unseen tokens: zqxjkv"]:
|
|
ids = tokenizer.encode(text)
|
|
assert tokenizer.decode(ids) == text
|
|
|
|
|
|
def test_special_tokens(tokenizer):
|
|
# all special tokens encode to a unique single id
|
|
ids = [tokenizer.encode_special(t) for t in SPECIAL_TOKENS]
|
|
assert len(set(ids)) == len(SPECIAL_TOKENS)
|
|
assert tokenizer.get_bos_token_id() == tokenizer.encode_special("<|bos|>")
|
|
# specials are NOT special-cased in ordinary text encoding
|
|
ids = tokenizer.encode("<|bos|>")
|
|
assert len(ids) > 1, "special token strings in plain text should not collapse to one token"
|
|
|
|
|
|
def test_encode_prepend_append(tokenizer):
|
|
bos = tokenizer.get_bos_token_id()
|
|
ids = tokenizer.encode("hello", prepend="<|bos|>", append="<|user_end|>")
|
|
assert ids[0] == bos
|
|
assert ids[-1] == tokenizer.encode_special("<|user_end|>")
|
|
|
|
|
|
def test_encode_batch(tokenizer):
|
|
texts = ["hello", "world"]
|
|
ids = tokenizer.encode(texts)
|
|
assert isinstance(ids, list) and len(ids) == 2
|
|
assert ids[0] == tokenizer.encode("hello")
|
|
|
|
|
|
def test_render_conversation_masks(tokenizer):
|
|
conversation = {"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
{"role": "assistant", "content": "hello!"},
|
|
{"role": "user", "content": "bye"},
|
|
{"role": "assistant", "content": "later"},
|
|
]}
|
|
ids, mask = tokenizer.render_conversation(conversation)
|
|
assert len(ids) == len(mask)
|
|
# first token is bos and is not supervised
|
|
assert ids[0] == tokenizer.get_bos_token_id() and mask[0] == 0
|
|
# supervised tokens are exactly: assistant content + assistant_end tokens
|
|
assistant_end = tokenizer.encode_special("<|assistant_end|>")
|
|
supervised_ids = [i for i, m in zip(ids, mask) if m == 1]
|
|
expected = tokenizer.encode("hello!") + [assistant_end] + tokenizer.encode("later") + [assistant_end]
|
|
assert supervised_ids == expected
|
|
# user content tokens are never supervised
|
|
user_start = tokenizer.encode_special("<|user_start|>")
|
|
assert all(m == 0 for i, m in zip(ids, mask) if i == user_start)
|
|
|
|
|
|
def test_render_conversation_system_message_merged(tokenizer):
|
|
without_system = {"messages": [
|
|
{"role": "user", "content": "sys prompt\n\nhi"},
|
|
{"role": "assistant", "content": "yo"},
|
|
]}
|
|
with_system = {"messages": [
|
|
{"role": "system", "content": "sys prompt"},
|
|
{"role": "user", "content": "hi"},
|
|
{"role": "assistant", "content": "yo"},
|
|
]}
|
|
assert tokenizer.render_conversation(with_system) == tokenizer.render_conversation(without_system)
|
|
|
|
|
|
def test_render_conversation_tool_parts(tokenizer):
|
|
# python tool calls are supervised, python outputs (come from the interpreter) are not
|
|
conversation = {"messages": [
|
|
{"role": "user", "content": "add"},
|
|
{"role": "assistant", "content": [
|
|
{"type": "text", "text": "sure"},
|
|
{"type": "python", "text": "1+1"},
|
|
{"type": "python_output", "text": "2"},
|
|
{"type": "text", "text": "it is 2"},
|
|
]},
|
|
]}
|
|
ids, mask = tokenizer.render_conversation(conversation)
|
|
python_start = tokenizer.encode_special("<|python_start|>")
|
|
output_start = tokenizer.encode_special("<|output_start|>")
|
|
output_end = tokenizer.encode_special("<|output_end|>")
|
|
# the tool call and its delimiters are supervised
|
|
assert mask[ids.index(python_start)] == 1
|
|
# the interpreter output and its delimiters are not
|
|
start, end = ids.index(output_start), ids.index(output_end)
|
|
assert all(m == 0 for m in mask[start:end + 1])
|
|
|
|
|
|
def test_render_conversation_truncation(tokenizer):
|
|
conversation = {"messages": [
|
|
{"role": "user", "content": "hello " * 100},
|
|
{"role": "assistant", "content": "world " * 100},
|
|
]}
|
|
ids, mask = tokenizer.render_conversation(conversation, max_tokens=32)
|
|
assert len(ids) == 32 and len(mask) == 32
|
|
|
|
|
|
def test_render_for_completion(tokenizer):
|
|
conversation = {"messages": [
|
|
{"role": "user", "content": "hi"},
|
|
{"role": "assistant", "content": "this gets stripped"},
|
|
]}
|
|
ids = tokenizer.render_for_completion(conversation)
|
|
# ends with assistant_start, primed for a completion
|
|
assert ids[-1] == tokenizer.encode_special("<|assistant_start|>")
|
|
# the assistant response itself must not be present
|
|
stripped = tokenizer.encode("this gets stripped")
|
|
assert not any(ids[i:i + len(stripped)] == stripped for i in range(len(ids)))
|