Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix stop detections #1392

Merged
merged 7 commits into from
May 22, 2023
Merged

fix stop detections #1392

Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c083920
fix stop detections
mingfang May 20, 2023
dfde374
rename to partial_stop and add comment
mingfang May 20, 2023
101438e
refine partial_stop()
mingfang May 21, 2023
3c40b62
avoid using same partial_stop name for both flag and function
mingfang May 21, 2023
cc906f4
make partial_stop() global
mingfang May 21, 2023
fde31b8
default partially_stopped to False when stop_str is False
mingfang May 21, 2023
ffe382f
default partially_stopped to False
mingfang May 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 26 additions & 10 deletions fastchat/serve/inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,13 @@ def prepare_logits_processor(
return processor_list


def partial_stop(output, stop_str):
for i in range(0, min(len(output), len(stop_str))):
if stop_str.startswith(output[-i:]):
return True
return False


@torch.inference_mode()
def generate_stream(
model, tokenizer, params, device, context_len=2048, stream_interval=2
Expand Down Expand Up @@ -160,31 +167,40 @@ def generate_stream(
skip_special_tokens=True,
spaces_between_special_tokens=False,
)

if stop_str:
if isinstance(stop_str, str):
pos = output.rfind(stop_str, rfind_start)
if pos != -1:
output = output[:pos]
stopped = True
else:
partially_stopped = partial_stop(output, stop_str)
elif isinstance(stop_str, Iterable):
for each_stop in stop_str:
pos = output.rfind(each_stop, rfind_start)
if pos != -1:
output = output[:pos]
stopped = True
break
else:
partially_stopped = partial_stop(output, each_stop)
if partially_stopped:
break
else:
raise ValueError("Invalid stop field type.")

yield {
"text": output,
"usage": {
"prompt_tokens": input_echo_len,
"completion_tokens": i,
"total_tokens": input_echo_len + i,
},
"finish_reason": None,
}

# prevent yielding partial stop sequence
if not partially_stopped:
mingfang marked this conversation as resolved.
Show resolved Hide resolved
yield {
"text": output,
"usage": {
"prompt_tokens": input_echo_len,
"completion_tokens": i,
"total_tokens": input_echo_len + i,
},
"finish_reason": None,
}

if stopped:
break
Expand Down