nanochat / tests /test_tokenizer.py

Andrej Karpathy

add a few tests

f8a85a5 4 months ago

5.49 kB

"""
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)))

	"""
	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)))