Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-12-23 21:18:00 +01:00
Code Issues Packages Projects Releases 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
modules/RWKV.py
View File

@ -31,5 +31,14 @@ class RWKVModel:
result.model = pipeline
return result
def generate(self, context, **kwargs):
return self.model.generate(context, **kwargs)
def generate(self, context, token_count=20, temperature=1, top_p=1, alpha_frequency=0.25, alpha_presence=0.25, token_ban=[0], token_stop=[], callback=None):
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 = []
)
return self.model.generate(context, token_count=token_count, args=args, callback=callback)

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)