diff --git a/headroom/tokenizers/huggingface.py b/headroom/tokenizers/huggingface.py index 750971f33..8ace0a792 100644 --- a/headroom/tokenizers/huggingface.py +++ b/headroom/tokenizers/huggingface.py @@ -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: diff --git a/headroom/tokenizers/registry.py b/headroom/tokenizers/registry.py index eddce6fa6..d057474f9 100644 --- a/headroom/tokenizers/registry.py +++ b/headroom/tokenizers/registry.py @@ -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" diff --git a/tests/test_tokenizer_selection_coverage.py b/tests/test_tokenizer_selection_coverage.py new file mode 100644 index 000000000..2b7c2e045 --- /dev/null +++ b/tests/test_tokenizer_selection_coverage.py @@ -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"