Add back the max_updates_second parameter (#5937)

This commit is contained in:
oobabooga 2024-04-26 10:14:51 -03:00 committed by GitHub
parent 6761b5e7c6
commit 70845c76fb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 15 additions and 0 deletions

View file

@ -84,6 +84,10 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
state = copy.deepcopy(state)
state['stream'] = True
min_update_interval = 0
if state.get('max_updates_second', 0) > 0:
min_update_interval = 1 / state['max_updates_second']
# Generate
for reply in generate_func(question, original_question, seed, state, stopping_strings, is_chat=is_chat):
reply, stop_found = apply_stopping_strings(reply, all_stop_strings)
@ -101,7 +105,14 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
last_update = time.time()
yield reply
# Limit updates to avoid lag in the Gradio UI
# API updates are not limited
else:
if cur_time - last_update > min_update_interval:
last_update = cur_time
yield reply
yield reply
if stop_found or (state['max_tokens_second'] > 0 and shared.stop_everything):