DeepInfer/tests/serve/test_prompt_tokenizer.py

"""serve/prompt_tokenizer.py: every backend must give tokenizer.encode's ids (CPU only, needs the checkpoint's tokenizer)."""

from __future__ import annotations

import copy
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

import pytest

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "serve"))
from prompt_tokenizer import PromptTokenizer  # noqa: E402

MODEL = os.environ.get("DEEPINFER_HF_MODEL_DIR")
pytestmark = pytest.mark.skipif(not MODEL or not Path(MODEL, "tokenizer.json").exists(), reason="DEEPINFER_HF_MODEL_DIR not set")


@pytest.fixture(scope="module")
def tokenizer() -> Any:
    from transformers import AutoTokenizer

    return AutoTokenizer.from_pretrained(MODEL)


@pytest.fixture(scope="module")
def turns() -> list[str]:
    """A growing tool-use conversation rendered with the checkpoint's encode_messages, one prompt per turn."""
    sys.path.insert(0, str(Path(MODEL, "encoding")))
    from encoding import encode_messages

    files = sorted((REPO_ROOT / "serve").glob("*.py")) + sorted((REPO_ROOT / "deepinfer" / "engine").glob("*.py"))[:8]
    tools = [
        {
            "type": "function",
            "function": {"name": "read_file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}},
        }
    ]
    msgs = [
        {"role": "system", "content": "You are a coding agent. 你是一个编程助手。", "tools": tools},
        {"role": "user", "content": "Summarise serve/. 12345"},
    ]
    out = []
    for i, f in enumerate(files):
        call = {"type": "function", "function": {"name": "read_file", "arguments": json.dumps({"path": str(f.name)})}}
        msgs.append({"role": "assistant", "content": f"Step {i}: reading.", "tool_calls": [call]})
        msgs.append({"role": "tool", "content": f.read_text()})
        out.append(encode_messages(copy.deepcopy(msgs), thinking_mode="chat"))
    return out


def available(tokenizer: Any) -> list[str]:
    names = ["segment", "hf"]
    try:
        import fastokens  # noqa: F401

        names.insert(0, "fastokens")
    except ImportError:
        pass
    return names


def test_backends_match_hf(tokenizer: Any, turns: list[str]) -> None:
    for name in available(tokenizer):
        pt = PromptTokenizer(tokenizer, backend=name, verify_rate=0)
        assert pt.name == name
        for prompt in turns:
            assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False), name


def test_auto_picks_a_fast_backend(tokenizer: Any) -> None:
    assert PromptTokenizer(tokenizer, verify_rate=0).name in ("fastokens", "segment")


def test_segment_cache_reuses_earlier_turns(tokenizer: Any, turns: list[str]) -> None:
    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0)
    hf_s = seg_s = 0.0
    for prompt in turns:
        t0 = time.perf_counter()
        ref = tokenizer.encode(prompt, add_special_tokens=False)
        t1 = time.perf_counter()
        assert pt.encode(prompt) == ref
        hf_s, seg_s = hf_s + t1 - t0, seg_s + time.perf_counter() - t1
    assert seg_s < hf_s


def test_segment_eviction_stays_exact(tokenizer: Any, turns: list[str]) -> None:
    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0, cache_bytes=4096)
    for prompt in turns + turns[::-1]:
        assert pt.encode(prompt) == tokenizer.encode(prompt, add_special_tokens=False)


def test_segment_with_strip_tokens(tokenizer: Any) -> None:
    """Added tokens that swallow neighbouring whitespace (lstrip/rstrip, as Phi-3 and others use) stay exact."""
    from tokenizers import AddedToken

    tok = copy.deepcopy(tokenizer)
    tok.add_tokens([AddedToken("<|l|>", lstrip=True, normalized=False), AddedToken("<|r|>", rstrip=True, normalized=False)])
    pt = PromptTokenizer(tok, backend="segment", verify_rate=0)
    for text in ["a  <|l|>  b", "a<|r|> \n\t b", " <|l|><|r|> ", "x　<|l|> y<|r|> z", "<|r|>   ", "   <|l|>"]:
        assert pt.encode(text) == tok.encode(text, add_special_tokens=False), repr(text)


def test_runtime_mismatch_switches_to_hf(tokenizer: Any) -> None:
    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=1.0)
    pt._encode = lambda text: [0]  # a backend gone wrong after startup
    pt.encode("hello world")
    for _ in range(200):
        if pt.name == "hf":
            break
        time.sleep(0.01)
    assert pt.name == "hf" and pt.stats["verify_mismatches"] == 1
    assert pt.encode("hello world") == tokenizer.encode("hello world", add_special_tokens=False)


def test_backend_error_falls_back_to_hf(tokenizer: Any) -> None:
    pt = PromptTokenizer(tokenizer, backend="segment", verify_rate=0)

    def broken(text: str) -> list[int]:
        raise RuntimeError("regex limit")

    pt._encode = broken
    assert pt.encode("hello") == tokenizer.encode("hello", add_special_tokens=False)
    assert pt.stats["errors"] == 1


def test_forced_unavailable_backend_raises(tokenizer: Any, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setitem(sys.modules, "fastokens", None)
    monkeypatch.setitem(sys.modules, "fastokens._native", None)
    with pytest.raises(RuntimeError, match="fastokens"):
        PromptTokenizer(tokenizer, backend="fastokens")
    assert PromptTokenizer(tokenizer, verify_rate=0).name == "segment"