Commit graph

1 commit

Author SHA1 Message Date
Tejas Chopra
0ed306b22b
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)
2026-08-03 22:28:06 -07:00