Fix beam search eos (vllm-project#9627)

cooleel · Oct 28, 2024 · feb92fb · feb92fb
1 parent 32176fe
commit feb92fb
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py
@@ -140,7 +140,12 @@ async def beam_search(
         best_beams = sorted_completed[:beam_width]
 
         for beam in best_beams:
-            beam.text = tokenizer.decode(beam.tokens[tokenized_length:])
+            if (beam.tokens[-1] == tokenizer.eos_token_id and not ignore_eos):
+                # Skip the eos token in the text.
+                tokens = beam.tokens[tokenized_length:-1]
+            else:
+                tokens = beam.tokens[tokenized_length:]
+            beam.text = tokenizer.decode(tokens)
 
         beam_search_output = RequestOutput(
             request_id=request_id,