Add a flag to include stop string in output text (#1976)

This commit is contained in:
Yunfeng Bai
2023-12-15 00:45:58 -08:00
committed by GitHub
parent 614856da25
commit c06170cc8e
2 changed files with 32 additions and 24 deletions

View File

@@ -682,9 +682,10 @@ class LLMEngine:
"""Stop the finished sequences."""
for stop_str in sampling_params.stop:
if seq.output_text.endswith(stop_str):
# Truncate the output text so that the stop string is
# not included in the output.
seq.output_text = seq.output_text[:-len(stop_str)]
if not sampling_params.include_stop_str_in_output:
# Truncate the output text so that the stop string is
# not included in the output.
seq.output_text = seq.output_text[:-len(stop_str)]
seq.status = SequenceStatus.FINISHED_STOPPED
return
if seq.get_last_token_id() in sampling_params.stop_token_ids: