File size: 5,987 Bytes
dfb775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Canonical tokenizer wrapper (spec Rule 3.1).
The model's own `tokenizer.json` loaded via the HuggingFace `tokenizers`
(Rust) library is the authoritative token counter. Every reported token
figure carries the tokenizer revision hash, so cross-run comparisons
remain honest when tokenizer vocabularies differ (Qwen3.5's 248k vs
Qwen2.5's 152k vs Llama 3's 128k vocabulary).
The two cross-checks that matter operationally:
1. **Scaffold accounting (Rule 3.2)** — the ChatML wrapping adds 10-25
tokens per turn. `count_with_scaffold` reports both inclusive and
content-only counts so a "tok/s" figure can be disambiguated.
2. **Bytes-per-second (Rule 3.3)** — Qwen3.5 compresses English to ~3.5
chars/tok where Llama 3 sits at ~4.0; tokenizer-relative throughput
alone is misleading. `bytes_per_decoded_token` is the cross-vocabulary
honest answer.
Lazy import: this module loads without `tokenizers` installed; callers
get a clear RuntimeError pointing at `uv sync --extra ml`.
"""
from __future__ import annotations
import hashlib
from functools import lru_cache
from pathlib import Path
from typing import Any
def _require_tokenizers() -> Any:
try:
import tokenizers # type: ignore
except ImportError as exc:
msg = (
"MEI canonical tokenizer requires the `tokenizers` library — "
"run `uv sync --extra ml`."
)
raise RuntimeError(msg) from exc
return tokenizers
@lru_cache(maxsize=8)
def _load_tokenizer(model_id_or_path: str) -> Any:
"""Load Tokenizer.from_file(<dir>/tokenizer.json) or from_pretrained.
Cached so repeated count calls don't re-parse the JSON. The cache key
is the input string — callers can either pass a local path or an HF
Hub repo ID; both resolve through the same code path.
"""
tok_mod = _require_tokenizers()
path = Path(model_id_or_path).expanduser()
if path.is_dir() and (path / "tokenizer.json").exists():
return tok_mod.Tokenizer.from_file(str(path / "tokenizer.json"))
if path.is_file() and path.name == "tokenizer.json":
return tok_mod.Tokenizer.from_file(str(path))
# Fall back to HF Hub lookup via the transformers convenience layer if
# available; otherwise raise so the caller knows the lookup failed.
try:
return tok_mod.Tokenizer.from_pretrained(model_id_or_path)
except Exception as exc:
msg = (
f"could not load tokenizer for {model_id_or_path!r}: "
f"neither a local tokenizer.json nor an HF Hub repo could be "
f"resolved ({type(exc).__name__}: {exc})"
)
raise RuntimeError(msg) from exc
def tokenizer_revision(model_id_or_path: str) -> str:
"""Stable hash over the tokenizer.json bytes (spec Rule 3.1).
Every reported token figure carries this hash so a measurement run
can be replayed and cross-checked against the exact tokenizer
revision it consumed.
"""
path = Path(model_id_or_path).expanduser()
tj = path / "tokenizer.json" if path.is_dir() else path
if tj.exists() and tj.is_file():
digest = hashlib.blake2b(tj.read_bytes(), digest_size=16).hexdigest()
return f"local:{digest}"
# For HF Hub references we can't BLAKE the local file (it isn't
# materialized yet at this layer). Surface the repo identity instead.
return f"hub:{model_id_or_path}"
def encode_with_revision(
text: str, model_id_or_path: str, *, add_special_tokens: bool = False,
) -> tuple[list[int], str]:
"""Return (token_ids, tokenizer_revision_hash). Rule 3.1."""
tok = _load_tokenizer(model_id_or_path)
encoding = tok.encode(text, add_special_tokens=add_special_tokens)
return list(encoding.ids), tokenizer_revision(model_id_or_path)
def _chatml_render(messages: list[dict[str, str]]) -> str:
"""Render a ChatML conversation the way Qwen tokenizers expect.
The MEI counts every token the model processes, *including* scaffold
(Rule 3.2). We render explicitly here rather than delegating to a
tokenizer's `apply_chat_template` because we want byte-stable output
regardless of template revisions — the count is what matters.
"""
parts: list[str] = []
for m in messages:
role = m.get("role", "user")
content = m.get("content", "")
parts.append(f"<|im_start|>{role}\n{content}<|im_end|>\n")
# Trailing assistant prompt — the model's generation prefix.
parts.append("<|im_start|>assistant\n")
return "".join(parts)
def count_with_scaffold(
messages: list[dict[str, str]],
model_id_or_path: str,
*,
include_scaffold: bool = True,
) -> int:
"""Token count for a chat-style payload (Rule 3.2).
When `include_scaffold=True` (the operationally-honest default), the
ChatML `<|im_start|>`, `<|im_end|>` markers and role headers are
counted — matches what the model actually processes. When False, the
count strips structural overhead — suitable for user-facing
reporting only, never for tok/s denominators.
"""
if include_scaffold:
rendered = _chatml_render(messages)
else:
rendered = "\n".join(m.get("content", "") for m in messages)
ids, _ = encode_with_revision(rendered, model_id_or_path)
return len(ids)
def bytes_per_decoded_token(
token_ids: list[int], model_id_or_path: str,
) -> float:
"""UTF-8 bytes per decoded token (Rule 3.3).
Cross-tokenizer-invariant cross-check. A Qwen 100 tok/s and a Llama
100 tok/s produce different amounts of useful output; this number is
how MEI keeps cross-vocab comparisons honest. Returns 0.0 for empty
token lists.
"""
if not token_ids:
return 0.0
tok = _load_tokenizer(model_id_or_path)
decoded = tok.decode(token_ids)
return len(decoded.encode("utf-8")) / len(token_ids)
__all__ = [
"bytes_per_decoded_token",
"count_with_scaffold",
"encode_with_revision",
"tokenizer_revision",
]
|