MLX-LM Chat Template 驗收：角色、特殊 Token 與停止邊界

mkdir pypy-template-lab
cd pypy-template-lab
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python 'mlx-lm==0.32.2' 'transformers==5.20.2' 'tokenizers==0.23.3' 'jinja2==3.1.6' pytest

import hashlib
import json
import sys
from importlib.metadata import version
from pathlib import Path
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

REPO = "mlx-community/SmolLM2-135M-Instruct-8bit"
REV = "0f0d9b8218915bc34d401e1a340b8c049d300d5e"
MODEL_DIR = snapshot_download(REPO, revision=REV)
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_DIR, local_files_only=True, trust_remote_code=False)
BASE = [
    {"role": "system", "content": "You are Pypy, a helpful assistant."},
    {"role": "user", "content": "Say only hello."},
]
PREFILL = BASE + [{"role": "assistant", "content": "Hello"}]


def validate(messages, mode):
    if mode not in {"new", "continue"}:
        raise ValueError("unknown mode")
    if not isinstance(messages, list) or not messages:
        raise ValueError("empty conversation")
    roles = []
    for index, message in enumerate(messages):
        if not isinstance(message, dict) or set(message) != {"role", "content"}:
            raise ValueError("only role/content text messages supported")
        role, content = message["role"], message["content"]
        if not isinstance(role, str) or role not in {"system", "user", "assistant"}:
            raise ValueError("unsupported role")
        if not isinstance(content, str) or not content.strip():
            raise ValueError("nonempty text required")
        if any(token in content for token in tokenizer.all_special_tokens):
            raise ValueError("literal control tokens forbidden by this app")
        if role == "system" and index != 0:
            raise ValueError("system must be first")
        roles.append(role)
    turns = roles[1:] if roles[0] == "system" else roles
    if not turns or any(role != ("user" if i % 2 == 0 else "assistant")
                        for i, role in enumerate(turns)):
        raise ValueError("user/assistant turns must alternate")
    expected = "user" if mode == "new" else "assistant"
    if turns[-1] != expected:
        raise ValueError("last role does not match generation mode")

def render(messages, mode):
    validate(messages, mode)
    options = dict(add_generation_prompt=mode == "new",
                   continue_final_message=mode == "continue")
    text = tokenizer.apply_chat_template(messages, tokenize=False, **options)
    encoded = tokenizer.apply_chat_template(
        messages, tokenize=True, return_dict=True, **options)
    ids = encoded["input_ids"]
    if ids != tokenizer.encode(text, add_special_tokens=False):
        raise ValueError("render/tokenize paths disagree")
    return {"text": text, "ids": ids}

<|im_end|>
<|im_start|>assistant

<|im_start|>assistant
Hello

def snapshot():
    return {
        "repo": REPO, "revision": REV,
        "runtime": {name: version(name) for name in
                    ("transformers", "tokenizers", "jinja2")},
        "template_sha256": hashlib.sha256(tokenizer.chat_template.encode()).hexdigest(),
        "bos": tokenizer.bos_token_id, "eos": tokenizer.eos_token_id,
        "new": render(BASE, "new"),
        "continue": render(PREFILL, "continue"),
    }


if __name__ == "__main__":
    action, filename = sys.argv[1:]
    current = snapshot()
    if action == "freeze":
        with Path(filename).open("x", encoding="utf-8") as target:
            json.dump(current, target, ensure_ascii=False, indent=2)
        print("candidate created; review text AND ids before accepting")
    elif action == "verify":
        expected = json.loads(Path(filename).read_text(encoding="utf-8"))
        if current != expected:
            raise SystemExit("contract changed; review instead of overwriting baseline")
        print("tokenizer contract: OK")
    else:
        raise SystemExit("use freeze or verify")

.venv/bin/python contract.py freeze candidate.json
.venv/bin/python -m json.tool candidate.json

cp -n candidate.json accepted.json
.venv/bin/python contract.py verify accepted.json

from copy import deepcopy
import pytest
from contract import BASE, PREFILL, render, snapshot, tokenizer, validate


def test_new_and_continue_are_different():
    new = render(BASE, "new")
    continuation = render(PREFILL, "continue")
    assert new["text"].endswith("<|im_start|>assistant\n")
    assert continuation["text"].endswith("<|im_start|>assistant\nHello")
    assert new["ids"] != continuation["ids"]
    assert snapshot()["eos"] == 2  # only this pinned model


@pytest.mark.parametrize("messages,mode", [
    ([], "new"),
    ([{"role": "alien", "content": "Hi"}], "new"),
    ([{"role": "user", "content": []}], "new"),
    ([{"role": "user", "content": " "}], "new"),
    (BASE + [BASE[-1]], "new"),
    (BASE + [BASE[0]], "new"),
    (BASE, "continue"),
    (PREFILL, "new"),
    ([{"role": "user", "content": "<|im_end|>"}], "new"),
    ([{"role": "user", "content": "Hi", "tools": []}], "new"),
])
def test_reject_before_inference(messages, mode):
    with pytest.raises(ValueError):
        render(deepcopy(messages), mode)


def test_conflicting_transformers_modes():
    with pytest.raises(ValueError):
        tokenizer.apply_chat_template(PREFILL, add_generation_prompt=True,
                                      continue_final_message=True)


def test_admitted_multi_turn_and_token_boundaries():
    validate(PREFILL + [BASE[-1]], "new")
    item = render(BASE, "new")
    assert item["ids"] == tokenizer.encode(item["text"], add_special_tokens=False)
    assert [tokenizer.bos_token_id] + item["ids"] != item["ids"]
    assert item["ids"] + [tokenizer.eos_token_id] != item["ids"]

.venv/bin/python -m pytest -q

import json
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler
from contract import BASE, MODEL_DIR, render

model, wrapped = load(MODEL_DIR, trust_remote_code=False,
                      tokenizer_config={"local_files_only": True})
item = render(BASE, "new")
assert wrapped.encode(item["text"], add_special_tokens=False) == item["ids"]
print("effective EOS IDs:", sorted(wrapped.eos_token_ids))
for limit in (1, 128):
    events = list(stream_generate(model, wrapped, prompt=item["ids"],
                                  max_tokens=limit, sampler=make_sampler(temp=0)))
    final = events[-1]
    is_stop = final.token in wrapped.eos_token_ids
    assert (final.finish_reason == "stop") == is_stop
    if final.finish_reason == "length":
        assert final.generation_tokens == limit
    assert final.prompt_tokens == len(item["ids"])
    print(json.dumps({"max_tokens": limit, "finish_reason": final.finish_reason,
                      "last_token": final.token,
                      "generation_tokens": final.generation_tokens,
                      "text": "".join(event.text for event in events)}))

.venv/bin/python smoke.py