Skip to content

fix(analyze): count prompt characters, not dict length, in count_tokens - #2031

Open
manunicholasjacob wants to merge 1 commit into
NVIDIA:mainfrom
manunicholasjacob:fix/count-tokens-prompt-length
Open

fix(analyze): count prompt characters, not dict length, in count_tokens#2031
manunicholasjacob wants to merge 1 commit into
NVIDIA:mainfrom
manunicholasjacob:fix/count-tokens-prompt-length

Conversation

@manunicholasjacob

Copy link
Copy Markdown
Contributor

Summary

count_tokens reports the wrong input character count. r["prompt"] is a serialised
Turn, which has exactly two fields (role, content), so len(r["prompt"]) is always
2 regardless of how long the prompt is.

The module docstring describes counting "the number of characters sent and received", and
the function already unwraps dict-shaped outputs via a nested _to_text helper - the
prompt side never got the same treatment.

Effect

A report with 100 attempts, 4000-character prompts and run.generations: 5 prints
Input chars: 1000 instead of 2,000,000. The figure is independent of prompt length.

What this changes

Promotes _to_text to module level, makes it recursive so it can reach
prompt.content.text, and uses it for the prompt as well as the outputs. Reports with a
plain-string prompt still work.

Testing

Adds tests/analyze/test_count_tokens.py covering a nested prompt and a legacy
string prompt. The nested case fails on main and passes with this change.

count_tokens called len() directly on r["prompt"], which is a dict for reports
written by generators that store structured prompts, so the token count was the
number of dict keys rather than the prompt length. The outputs path already had
a local _to_text helper for exactly this; hoist it to module scope and use it
for the prompt as well.

Adds a regression test that fails on the old code path.

Signed-off-by: Manu Nicholas Jacob <manunicholasjacob@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant