mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add a generate() function for RWKV

Browse source
This commit is contained in:
oobabooga 2023-03-01 12:16:11 -03:00
parent 659bb76722
commit e735806c51
2 changed files with 13 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

10
modules/text_generation.py
View file

@ -85,19 +85,13 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
torch.cuda.empty_cache()
if shared.is_RWKV:
args = PIPELINE_ARGS(temperature = temperature, top_p = top_p,
alpha_frequency = 0.25, # Frequency Penalty (as in GPT-3)
alpha_presence = 0.25, # Presence Penalty (as in GPT-3)
token_ban = [0], # ban the generation of some tokens
token_stop = []) # stop generation whenever you see any token here
if shared.args.no_stream:
reply = question + shared.model.generate(question, token_count=max_new_tokens, args=args, callback=None)
reply = question + shared.model.generate(question, token_count=max_new_tokens, temperature=temperature)
yield formatted_outputs(reply, None)
return formatted_outputs(reply, None)
else:
for i in range(max_new_tokens//8):
reply = question + shared.model.generate(question, token_count=8, args=args, callback=None)
reply = question + shared.model.generate(question, token_count=8, temperature=temperature)
yield formatted_outputs(reply, None)
question = reply
return formatted_outputs(reply, None)