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

RWKV support prototype

Browse source
This commit is contained in:
oobabooga 2023-02-27 23:03:35 -03:00
parent 021bd55886
commit ebc64a408c
3 changed files with 42 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

14
modules/text_generation.py
View file

@ -6,6 +6,7 @@ import numpy as np
import torch
import transformers
from tqdm import tqdm
from rwkv.utils import PIPELINE, PIPELINE_ARGS
import modules.shared as shared
from modules.extensions import apply_extensions
@ -80,6 +81,19 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
if not shared.args.cpu:
torch.cuda.empty_cache()
if shared.is_RWKV:
def my_print(s):
print(s, end='', flush=True)
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
reply = question + shared.model.generate(question, token_count=max_new_tokens, args=args, callback=None)
print(formatted_outputs(reply, None))
yield formatted_outputs(reply, None)
return formatted_outputs(reply, None)
original_question = question
if not (shared.args.chat or shared.args.cai_chat):
question = apply_extensions(question, "input")