mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(tokenizers): count HuggingFace chat templates, and resolve gpt-5 / gateway-wrapped names (#2758)
## Description Three tokenizer-selection defects, all measured against real counters on identical text. ### 1. HuggingFace-routed models counted a whole conversation as **2 tokens** `transformers >= 5` defaults `apply_chat_template(tokenize=True)` to `return_dict=True` and returns a `BatchEncoding`, so `len(formatted)` counted **dict keys** — `input_ids`, `attention_mask` — instead of tokens. ```text Qwen2.5-72B, one 6,000-char message before: count_messages = 2 count_message = -1 after : count_messages = 1020 count_message = 1017 true : ~1003 ``` `count_message` goes negative because `BaseTokenizer` subtracts a 3-token reply overhead from it. A **~99.8% undercount** on every HF-routed family whose resolved tokenizer carries a chat template — llama, qwen, deepseek, phi, yi, falcon, starcoder. `pyproject.toml` pins `transformers>=5.5.0,<6.0`, so the affected version is the only installable one, and nothing covered `count_messages`. It hid behind a second bug while I reproduced it: `DeepSeek-V3` mis-resolves to `deepseek-llm-7b-base` (a 2023 model with **no** chat template), which falls back to the estimator and looks fine. That mis-resolution is left for a follow-up. ### 2. The current OpenAI flagships had no pattern `MODEL_PATTERNS` stopped at `^gpt-4` / `^o1` / `^o3`: ```text gpt-5, gpt-5.1, gpt-5-mini, gpt-5.1-codex, o4-mini -> EstimatingTokenCounter ``` Deviation vs the correct `o200k` encoding: **+20% English, -33% JSON, -44% logs.** ### 3. Every pattern is `^`-anchored, so gateway-wrapped ids matched nothing ```text bedrock/anthropic.claude-3-5-sonnet -> EstimatingTokenCounter vertex_ai/claude-sonnet-4-6 -> EstimatingTokenCounter openrouter/anthropic/claude-sonnet-4-6 -> EstimatingTokenCounter us.anthropic.claude-sonnet-4-6-v1:0 -> EstimatingTokenCounter azure/gpt-4o -> EstimatingTokenCounter ``` Deviation: **+15% English, -33% JSON, -38% logs.** Not hypothetical — `handlers/openai.py` already documents that LiteLLM's `headroom` guardrail passes exactly these forms. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `tokenizers/huggingface.py` — pass `return_dict=False` to `apply_chat_template`. - `tokenizers/registry.py` — add `^gpt-5` and `^o4` to `MODEL_PATTERNS`. - `tokenizers/registry.py` — new `_name_candidates()`; `_detect_backend` now tries progressively-unwrapped forms: path segments stripped left-to-right, then Bedrock's dotted `[region.]vendor.model`. **Why candidates rather than rewriting the name:** the full name is candidate 0, so no currently-correct resolution can move, and an unknown alias still falls back to estimation rather than matching by accident. The estimator is a legitimate *fallback*; the bug was reaching it when a real tokenizer for that family exists. ## Testing - [x] Unit tests pass - [x] Linting passes (`ruff check` + `format --check`, pinned 0.15.17) - [x] New tests added - [x] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/ tests/test_tokenizer_selection_coverage.py --exclude headroom/dashboard/templates All checks passed! $ pytest tests/test_tokenizer_selection_coverage.py -q 20 passed in 0.60s $ pytest tests/test_huggingface_tokenizer_timeout.py tests/test_tokenizers/ -q this branch: 12 passed clean upstream/main: 12 passed <- no regression ``` 20 new tests cover all three defects **plus** the no-regression cases: bare names unchanged, unknown aliases still estimated, wrapped Gemini matching its bare form exactly, and candidate ordering/dedup. ## Real Behavior Proof - **Environment:** macOS 26.4 arm64, isolated worktree off `upstream/main`. The HF measurement used a real `transformers 5.14.1` with `Qwen/Qwen2.5-72B` from the local HF cache. **After the fix, resolution across every form a gateway realistically sends:** ```text gpt-4o TiktokenCounter gpt-5 TiktokenCounter <- was Estimating gpt-5.1 TiktokenCounter <- was Estimating o3-mini TiktokenCounter o4-mini TiktokenCounter <- was Estimating claude-sonnet-4-6 TiktokenCounter bedrock/anthropic.claude-3-5-sonnet TiktokenCounter <- was Estimating anthropic.claude-3-5-sonnet-20241022-v2:0 TiktokenCounter <- was Estimating us.anthropic.claude-sonnet-4-6-v1:0 TiktokenCounter <- was Estimating vertex_ai/claude-sonnet-4-6 TiktokenCounter <- was Estimating openrouter/anthropic/claude-sonnet-4-6 TiktokenCounter <- was Estimating azure/gpt-4o TiktokenCounter <- was Estimating vertex_ai/gemini-2.5-pro EstimatingTokenCounter (google backend, correct) groq/llama-3.3-70b-versatile HuggingFaceTokenizer <- was Estimating my-gateway/big-model EstimatingTokenCounter (correct fallback) ``` `vertex_ai/gemini-2.5-pro` and `gemini-2.5-pro` return **identical** counts (600 on the same input), confirming the prefix strip reaches the google backend rather than the generic fallback. - **Not fully tested locally:** `tests/test_evals_cjk_tokenization.py` cannot collect in this env — `ModuleNotFoundError: headroom._core`, the compiled Rust extension this machine can't currently build. Identical on baseline, so CI is the check there. It is CJK-related and this PR changes encoding selection for `gpt-5`/`o4`/wrapped names, so it's the suite most worth watching. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] I did **not** edit `CHANGELOG.md` ## Known follow-ups, deliberately not here - `DeepSeek-V3` → `deepseek-llm-7b-base`, `Qwen/Qwen2.5-72B` → `Qwen/Qwen-7B`: `get_tokenizer_name` prefix-matches against the whole string including the org segment, and has no version boundary. - `get_encoding_for_model` is case-sensitive while `_detect_backend` lowercases, so `GPT-4O` gets `cl100k` (+38.9% on CJK). - `providers/openai.py` has a second, divergent encoding resolver — it disagrees with `tokenizers/` on `gpt-4.1`, `gpt-5`, `text-embedding-3-large`, `davinci`. - Provider counters price most modern content blocks at literally zero (`thinking`, `document`, `mcp_tool_result`, and OpenAI's own `output_text`/`refusal`). 🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Tejas Chopra
GitHub
parent
184146b688
commit
0ed306b22b
3 changed files with 162 additions and 4 deletions
|
|
@ -325,11 +325,22 @@ class HuggingFaceTokenizer(BaseTokenizer):
|
|||
# Try to use chat template for accurate counting
|
||||
if hasattr(self.tokenizer, "apply_chat_template"):
|
||||
try:
|
||||
# Apply chat template and count
|
||||
# ``return_dict=False`` is load-bearing. transformers >= 5 defaults
|
||||
# ``apply_chat_template(tokenize=True)`` to ``return_dict=True``,
|
||||
# which hands back a BatchEncoding — so ``len(formatted)`` counted
|
||||
# DICT KEYS (2: input_ids, attention_mask) instead of tokens.
|
||||
# Measured on Qwen2.5-72B, a 6,000-char message: count_messages
|
||||
# returned 2 and count_message returned -1 (base subtracts a
|
||||
# 3-token reply overhead), against a true 1,003 tokens. That is a
|
||||
# ~99.8% undercount on every HF-routed family whose resolved
|
||||
# tokenizer carries a chat template — llama, qwen, deepseek, phi,
|
||||
# yi, falcon, starcoder. pyproject pins transformers>=5.5.0,<6.0,
|
||||
# so the affected version is the only installable one.
|
||||
formatted = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
return_dict=False,
|
||||
)
|
||||
return len(formatted)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -24,11 +24,13 @@ logger = logging.getLogger(__name__)
|
|||
# Order matters - more specific patterns first
|
||||
MODEL_PATTERNS: list[tuple[str, str]] = [
|
||||
# OpenAI models -> tiktoken
|
||||
(r"^gpt-5", "tiktoken"),
|
||||
(r"^gpt-4o", "tiktoken"),
|
||||
(r"^gpt-4", "tiktoken"),
|
||||
(r"^gpt-3\.5", "tiktoken"),
|
||||
(r"^o1", "tiktoken"),
|
||||
(r"^o3", "tiktoken"),
|
||||
(r"^o4", "tiktoken"),
|
||||
(r"^text-embedding", "tiktoken"),
|
||||
(r"^text-davinci", "tiktoken"),
|
||||
(r"^code-", "tiktoken"),
|
||||
|
|
@ -74,6 +76,42 @@ MODEL_PATTERNS: list[tuple[str, str]] = [
|
|||
]
|
||||
|
||||
|
||||
def _name_candidates(model_lower: str) -> tuple[str, ...]:
|
||||
"""Progressively-unwrapped forms of a model name, most specific first.
|
||||
|
||||
Every entry in :data:`MODEL_PATTERNS` is anchored with ``^``, which is right
|
||||
for a bare model id and wrong for the wrapped ids gateways actually send. A
|
||||
name like ``bedrock/anthropic.claude-sonnet-4-6-v1:0`` matched nothing and
|
||||
fell through to the char estimator instead of the Claude counter — measured
|
||||
deviation on identical text: +15% English, -33% JSON, -38% logs. Affected
|
||||
every ``bedrock/``, ``vertex_ai/``, ``openrouter/``, ``anthropic/``,
|
||||
``azure/``, ``groq/`` and ``litellm/`` form, plus Bedrock's bare
|
||||
``anthropic.claude-…`` and its ``us.``/``eu.``/``apac.`` region variants.
|
||||
|
||||
Yielding candidates rather than rewriting the name keeps the exact-match case
|
||||
first, so no currently-correct resolution can change.
|
||||
"""
|
||||
seen: list[str] = []
|
||||
|
||||
def add(name: str) -> None:
|
||||
if name and name not in seen:
|
||||
seen.append(name)
|
||||
|
||||
add(model_lower)
|
||||
# Strip provider path segments left-to-right: openrouter/anthropic/claude-x
|
||||
# yields anthropic/claude-x then claude-x.
|
||||
rest = model_lower
|
||||
while "/" in rest:
|
||||
rest = rest.split("/", 1)[1]
|
||||
add(rest)
|
||||
# Bedrock dotted ids: [region.]vendor.model
|
||||
for candidate in list(seen):
|
||||
parts = candidate.split(".")
|
||||
for i in range(1, len(parts)):
|
||||
add(".".join(parts[i:]))
|
||||
return tuple(seen)
|
||||
|
||||
|
||||
class TokenizerRegistry:
|
||||
"""Registry for tokenizer instances and factories.
|
||||
|
||||
|
|
@ -289,9 +327,10 @@ class TokenizerRegistry:
|
|||
"""
|
||||
model_lower = model.lower()
|
||||
|
||||
for pattern, backend in MODEL_PATTERNS:
|
||||
if re.match(pattern, model_lower):
|
||||
return backend
|
||||
for candidate in _name_candidates(model_lower):
|
||||
for pattern, backend in MODEL_PATTERNS:
|
||||
if re.match(pattern, candidate):
|
||||
return backend
|
||||
|
||||
# Default to estimation for unknown models
|
||||
return "estimation"
|
||||
|
|
|
|||
108
tests/test_tokenizer_selection_coverage.py
Normal file
108
tests/test_tokenizer_selection_coverage.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
"""Model names must resolve to the tokenizer their model actually uses.
|
||||
|
||||
Two selection gaps, both measured against real counters on identical text:
|
||||
|
||||
1. ``MODEL_PATTERNS`` stopped at ``^gpt-4``/``^o1``/``^o3``, so the current
|
||||
flagships — ``gpt-5``, ``gpt-5.1``, ``o4-mini`` — fell through to the char
|
||||
estimator. Deviation vs the correct o200k encoding: +20% English, -33% JSON,
|
||||
-44% logs.
|
||||
|
||||
2. Every pattern is ``^``-anchored, which is right for a bare model id and wrong
|
||||
for the wrapped ids gateways send. ``bedrock/anthropic.claude-3-5-sonnet``,
|
||||
``vertex_ai/claude-…``, ``openrouter/anthropic/claude-…``, ``azure/gpt-4o``
|
||||
and Bedrock's ``us.anthropic.claude-…`` all matched nothing. LiteLLM's
|
||||
``headroom`` guardrail passes exactly these forms.
|
||||
|
||||
The estimator is a legitimate FALLBACK; the bug is reaching it when a real
|
||||
tokenizer for that family exists.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.tokenizers.registry import _name_candidates
|
||||
|
||||
_TIKTOKEN = "TiktokenCounter"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"gpt-5",
|
||||
"gpt-5.1",
|
||||
"gpt-5-mini",
|
||||
"gpt-5.1-codex",
|
||||
"o4-mini",
|
||||
],
|
||||
)
|
||||
def test_current_openai_flagships_get_a_real_tokenizer(model: str) -> None:
|
||||
"""These fell to EstimatingTokenCounter before ^gpt-5 / ^o4 were added."""
|
||||
assert type(get_tokenizer(model)).__name__ == _TIKTOKEN
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
# gateway path prefixes
|
||||
"bedrock/anthropic.claude-3-5-sonnet",
|
||||
"vertex_ai/claude-sonnet-4-6",
|
||||
"openrouter/anthropic/claude-sonnet-4-6",
|
||||
"anthropic/claude-opus-4",
|
||||
"litellm/claude-sonnet-4-6",
|
||||
# Bedrock dotted ids, with and without a region segment
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"us.anthropic.claude-sonnet-4-6-v1:0",
|
||||
"eu.anthropic.claude-sonnet-4-6-v1:0",
|
||||
# OpenAI behind a gateway
|
||||
"azure/gpt-4o",
|
||||
"openrouter/openai/gpt-4o",
|
||||
],
|
||||
)
|
||||
def test_gateway_wrapped_names_resolve_like_their_bare_form(model: str) -> None:
|
||||
assert type(get_tokenizer(model)).__name__ == _TIKTOKEN
|
||||
|
||||
|
||||
def test_wrapped_gemini_matches_the_bare_form_exactly() -> None:
|
||||
"""Prefix stripping must reach the google backend, not the generic fallback."""
|
||||
text = "hello world " * 200
|
||||
assert get_tokenizer("vertex_ai/gemini-2.5-pro").count_text(text) == get_tokenizer(
|
||||
"gemini-2.5-pro"
|
||||
).count_text(text)
|
||||
|
||||
|
||||
def test_bare_names_are_unaffected() -> None:
|
||||
"""The exact-match candidate is tried first, so nothing already-correct moves."""
|
||||
for model, expected in (
|
||||
("gpt-4o", _TIKTOKEN),
|
||||
("gpt-3.5-turbo", _TIKTOKEN),
|
||||
("o1-preview", _TIKTOKEN),
|
||||
("o3-mini", _TIKTOKEN),
|
||||
("claude-sonnet-4-6", _TIKTOKEN),
|
||||
):
|
||||
assert type(get_tokenizer(model)).__name__ == expected, model
|
||||
|
||||
|
||||
def test_unknown_alias_still_falls_back_to_estimation() -> None:
|
||||
"""Prefix stripping must not invent a match for a genuinely unknown model."""
|
||||
assert type(get_tokenizer("my-gateway/big-model")).__name__ == "EstimatingTokenCounter"
|
||||
assert type(get_tokenizer("totally-unknown-xyz")).__name__ == "EstimatingTokenCounter"
|
||||
|
||||
|
||||
def test_name_candidates_orders_most_specific_first() -> None:
|
||||
"""The full name must be candidate 0 so exact registrations always win."""
|
||||
got = _name_candidates("openrouter/anthropic/claude-sonnet-4-6")
|
||||
assert got[0] == "openrouter/anthropic/claude-sonnet-4-6"
|
||||
assert "anthropic/claude-sonnet-4-6" in got
|
||||
assert "claude-sonnet-4-6" in got
|
||||
|
||||
dotted = _name_candidates("us.anthropic.claude-sonnet-4-6-v1:0")
|
||||
assert dotted[0] == "us.anthropic.claude-sonnet-4-6-v1:0"
|
||||
assert "claude-sonnet-4-6-v1:0" in dotted
|
||||
|
||||
|
||||
def test_name_candidates_is_deduplicated_and_finite() -> None:
|
||||
got = _name_candidates("a/b/c.d.e")
|
||||
assert len(got) == len(set(got))
|
||||
assert got[0] == "a/b/c.d.e"
|
||||
Loading…
Add table
Add a link
Reference in a new issue