From c6dfc3cdbeb4687a8313b82d399177e293fee196 Mon Sep 17 00:00:00 2001 From: xcnick Date: Wed, 12 Jul 2023 23:14:56 +0800 Subject: [PATCH] Fix handling of special tokens in decoding. (#418) --- vllm/engine/llm_engine.py | 5 +++-- vllm/transformers_utils/tokenizer.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index 6204191bbeec6..8936eda705014 100644 --- a/vllm/engine/llm_engine.py +++ b/vllm/engine/llm_engine.py @@ -276,8 +276,9 @@ class LLMEngine: seq.get_last_token_id(), skip_special_tokens=True, ) - seq.output_tokens.append(new_token) - seq.output_text = new_output_text + if new_token is not None: + seq.output_tokens.append(new_token) + seq.output_text = new_output_text def _stop_sequences(self, seq_groups: List[SequenceGroup]) -> None: """Stop the finished sequences.""" diff --git a/vllm/transformers_utils/tokenizer.py b/vllm/transformers_utils/tokenizer.py index 116a109a45414..4cb0fbed41ae8 100644 --- a/vllm/transformers_utils/tokenizer.py +++ b/vllm/transformers_utils/tokenizer.py @@ -80,6 +80,8 @@ def detokenize_incrementally( new_token: The new token as a string. output_text: The new output text as a string. """ + if skip_special_tokens and (new_token_id in tokenizer.all_special_ids): + return None, prev_output_tokens new_token = tokenizer.convert_ids_to_tokens( new_token_id, skip_special_tokens=skip_special_tokens) output_tokens = prev_output_tokens + [new_token] @@ -99,7 +101,7 @@ def detokenize_incrementally( sub_texts = [] current_sub_text = [] for token in output_tokens: - if skip_special_tokens and token in tokenizer.all_special_ids: + if skip_special_tokens and token in tokenizer.all_special_tokens: continue if token in tokenizer.added_tokens_encoder: if current_sub_text: