Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-22 08:07:56 +01:00
Code Issues Packages Projects Releases Wiki Activity

Fix a memory leak when text streaming is on

Browse Source
This commit is contained in:
oobabooga 2023-03-05 10:12:43 -03:00
parent 5492e2e9f8
commit 8e706df20e
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
modules/text_generation.py
View File

@ -81,11 +81,13 @@ def formatted_outputs(reply, model_name):
else:
return reply
def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=None, stopping_string=None):
def clear_torch_cache():
gc.collect()
if not shared.args.cpu:
torch.cuda.empty_cache()
def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typical_p, repetition_penalty, top_k, min_length, no_repeat_ngram_size, num_beams, penalty_alpha, length_penalty, early_stopping, eos_token=None, stopping_string=None):
clear_torch_cache()
t0 = time.time()
# These models are not part of Hugging Face, so we handle them
@ -98,6 +100,7 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
yield formatted_outputs(reply, shared.model_name)
else:
for i in tqdm(range(max_new_tokens//8+1)):
clear_torch_cache()
reply = shared.model.generate(question, token_count=8, temperature=temperature, top_p=top_p)
yield formatted_outputs(reply, shared.model_name)
question = reply
@ -183,6 +186,8 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
else:
yield formatted_outputs(original_question, shared.model_name)
for i in tqdm(range(max_new_tokens//8+1)):
clear_torch_cache()
with torch.no_grad():
output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0]
if shared.soft_prompt: