Fix: Correctly filter special tokens in benchmark_prefix_caching (#28615)

Signed-off-by: Di Wu <dw2761@nyu.edu>
2025-11-13 18:57:44 +08:00
parent 11ac9ddd03
commit e63fd44560
1 changed files with 1 additions and 1 deletions
--- a/benchmarks/benchmark_prefix_caching.py
+++ b/benchmarks/benchmark_prefix_caching.py
@@ -69,7 +69,7 @@ def sample_tokens(tokenizer: PreTrainedTokenizerBase, length: int) -> list[int]:

    # Remove the special tokens.
    return random.choices(
-        [v for k, v in vocab.items() if k not in all_special_ids],
+        [v for v in vocab.values() if v not in all_special_ids],
        k=length,
    )