Unverified Commit feb92fbe authored by Robert Shaw's avatar Robert Shaw Committed by GitHub
Browse files

Fix beam search eos (#9627)

parent 32176fee
......@@ -140,7 +140,12 @@ class EngineClient(ABC):
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,
......
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