Skip to content

perf: batch-tokenize contexts/continuations in _loglikelihood_tokens - #1386

Open
drmohanty-tech wants to merge 1 commit into
huggingface:mainfrom
drmohanty-tech:optimize-tokenization-732
Open

drmohanty-tech wants to merge 1 commit into
huggingface:mainfrom
drmohanty-tech:optimize-tokenization-732

Conversation

@drmohanty-tech

Copy link
Copy Markdown

What does this PR do?

Fixes #732.

_loglikelihood_tokens in TransformersModel currently tokenizes one document at a time in a Python loop, calling tok_encode_pair per document (which itself calls the tokenizer once for the context and once per continuation). For large benchmarks with many documents and/or many choices per document, this means a large number of small, separate tokenizer calls instead of one batched call, which is significantly slower than necessary.

This PR adds tok_encode_pair_batch to LightevalModel (in abstract_model.py): a pairwise-only batched equivalent of tok_encode_pair that:

  1. Applies the same trailing-context-space handling as tok_encode_pair, per document.
  2. Tokenizes all contexts in the batch with a single tokenizer call, and all continuations in the batch (flattened across documents) with a single tokenizer call.
  3. Redistributes the results back into the same per-document, per-choice nested structure that tok_encode_pair produces, including stripping a trailing EOS token from the context.

TransformersModel._loglikelihood_tokens is updated to call this new batched method instead of looping over the batch and calling tok_encode_pair per document.

Correctness

I verified that tok_encode_pair_batch produces byte-for-byte identical output to calling tok_encode_pair(..., pairwise=True) once per document, across multiple documents with varying numbers of choices and with/without trailing context spaces.

Tests

  • Added test_tok_encode_pair_batch_matches_per_document_tok_encode_pair, which asserts the batched output matches the per-document loop output exactly.
  • Added test_tok_encode_pair_batch_move_trailing_context_space, mirroring the existing test_tok_encode_pair_move_trailing_context_space test to confirm the move_trailing_context_space behavior is preserved in the batched path.
  • Existing tests in tests/unit/models/test_transformers_model.py (including test_loglikelihood_eval and test_loglikelihood_padded_tensors_shapes, which exercise _loglikelihood_tokens directly) pass unchanged.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case). N/A, not applicable.
  • Did you read the contributor guideline?
  • Was this discussed/approved via a Github issue? Please add a link to it if that's the case. See [BUG] Optimize tokenization #732.
  • Did you make sure to update the documentation with your changes? N/A, internal implementation detail with no user-facing API change.
  • Did you write any new necessary tests?

🤖 Generated with Claude Code

_loglikelihood_tokens previously called tok_encode_pair once per
document in the batch (and once more per continuation within each
document), each call hitting the tokenizer separately. This makes
tokenization scale linearly with per-item Python/call overhead
instead of letting the tokenizer batch across the whole batch.

Add tok_encode_pair_batch, a pairwise-only equivalent that tokenizes
all contexts and all continuations across the batch with a single
tokenizer call each, then redistributes results back per document.
Wire it into TransformersModel._loglikelihood_tokens in place of the
per-document loop.

Verified the new path produces identical token ids to the old
per-document loop, and added tests covering that equivalence plus
the move_trailing_context_space behavior in the batched path.

Fixes huggingface#732

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.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.

[BUG] Optimize tokenization

1 participant