Skip to content

Commit feb92fb

Browse files
Fix beam search eos (#9627)
1 parent 32176fe commit feb92fb

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

vllm/engine/protocol.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,12 @@ async def beam_search(
140140
best_beams = sorted_completed[:beam_width]
141141

142142
for beam in best_beams:
143-
beam.text = tokenizer.decode(beam.tokens[tokenized_length:])
143+
if (beam.tokens[-1] == tokenizer.eos_token_id and not ignore_eos):
144+
# Skip the eos token in the text.
145+
tokens = beam.tokens[tokenized_length:-1]
146+
else:
147+
tokens = beam.tokens[tokenized_length:]
148+
beam.text = tokenizer.decode(tokens)
144149

145150
beam_search_output = RequestOutput(
146151
request_id=request_id,

0 commit comments

Comments
 (0)