Skip to content

Commit

Permalink
Fix beam search eos (vllm-project#9627)
Browse files Browse the repository at this point in the history
Signed-off-by: Maxime Fournioux <55544262+mfournioux@users.noreply.github.com>
  • Loading branch information
robertgshaw2-redhat authored and mfournioux committed Nov 20, 2024
1 parent aecc1e0 commit b22e4f3
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion vllm/engine/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit b22e4f3

Please sign in to comment.