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]: ...@@ -69,7 +69,7 @@ def sample_tokens(tokenizer: PreTrainedTokenizerBase, length: int) -> list[int]:
# Remove the special tokens. # Remove the special tokens.
return random.choices( 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, 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