perf: batch-tokenize contexts/continuations in _loglikelihood_tokens - #1386
Open
drmohanty-tech wants to merge 1 commit into
Open
drmohanty-tech wants to merge 1 commit into
drmohanty-tech wants to merge 1 commit into
Conversation
_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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #732.
_loglikelihood_tokensinTransformersModelcurrently tokenizes one document at a time in a Python loop, callingtok_encode_pairper 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_batchtoLightevalModel(inabstract_model.py): a pairwise-only batched equivalent oftok_encode_pairthat:tok_encode_pair, per document.tok_encode_pairproduces, including stripping a trailing EOS token from the context.TransformersModel._loglikelihood_tokensis updated to call this new batched method instead of looping over the batch and callingtok_encode_pairper document.Correctness
I verified that
tok_encode_pair_batchproduces byte-for-byte identical output to callingtok_encode_pair(..., pairwise=True)once per document, across multiple documents with varying numbers of choices and with/without trailing context spaces.Tests
test_tok_encode_pair_batch_matches_per_document_tok_encode_pair, which asserts the batched output matches the per-document loop output exactly.test_tok_encode_pair_batch_move_trailing_context_space, mirroring the existingtest_tok_encode_pair_move_trailing_context_spacetest to confirm themove_trailing_context_spacebehavior is preserved in the batched path.tests/unit/models/test_transformers_model.py(includingtest_loglikelihood_evalandtest_loglikelihood_padded_tensors_shapes, which exercise_loglikelihood_tokensdirectly) pass unchanged.Before submitting
🤖 Generated with Claude Code