Unverified Commit e63fd445 authored by Di Wu's avatar Di Wu Committed by GitHub
Browse files

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


Signed-off-by: default avatarDi Wu <dw2761@nyu.edu>
parent 11ac9ddd
......@@ -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,
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment