import gradio as gr
import argparse
import os, gc, torch
from datetime import datetime
from huggingface_hub import hf_hub_download
import torch
# from pynvml import *
# nvmlInit()
# gpu_h = nvmlDeviceGetHandleByIndex(0)
ctx_limit = 4096
desc = f'''链接：太慢了？用Colab自己部署吧
ChatRWKVRWKV-LMRWKV pip package知乎教程
'''
parser = argparse.ArgumentParser(prog = 'ChatGal RWKV')
parser.add_argument('--share',action='store_true')
parser.add_argument('--ckpt',type=str,default="rwkv-loramerge-0426-v2-4096-epoch11.pth")
parser.add_argument('--model_path',type=str,default=None,help="local model path")
parser.add_argument('--lora', type=str, default=None, help='lora checkpoint path')
parser.add_argument('--lora_alpha', type=float, default=0, help='lora alpha')
parser.add_argument('--lora_layer_filter',type=str,default=None,help='layer filter. Default merge all layer. Example: "0.2*25-31"')
args = parser.parse_args()
os.environ["RWKV_JIT_ON"] = '1'
# from rwkv.model import RWKV
from rwkv_lora import RWKV
lora_kwargs = {
"lora":args.lora,
"lora_alpha":args.lora_alpha,
"lora_layer_filter":args.lora_layer_filter
}
if args.model_path:
model_path = args.model_path
else:
model_path = hf_hub_download(repo_id="Synthia/ChatGalRWKV", filename=args.ckpt)
# if 'ON_COLAB' in os.environ and os.environ['ON_COLAB'] == '1':
if torch.cuda.is_available() and torch.cuda.device_count()>0:
os.environ["RWKV_JIT_ON"] = '0'
os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)
model = RWKV(model=model_path, strategy='cuda bf16',**lora_kwargs)
else:
model = RWKV(model=model_path, strategy='cpu bf16',**lora_kwargs)
from utils import PIPELINE, PIPELINE_ARGS
pipeline = PIPELINE(model, "20B_tokenizer.json")
def infer(
ctx,
token_count=10,
temperature=0.7,
top_p=1.0,
top_k=50,
typical_p=1.0,
presencePenalty = 0.05,
countPenalty = 0.05,
):
args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p), top_k=int(top_k),typical_p=float(typical_p),
alpha_frequency = countPenalty,
alpha_presence = presencePenalty,
token_ban = [0], # ban the generation of some tokens
token_stop = []) # stop generation whenever you see any token here
# ctx = ctx.strip().split('\n')
# for c in range(len(ctx)):
# ctx[c] = ctx[c].strip().strip('\u3000').strip('\r')
# ctx = list(filter(lambda c: c != '', ctx))
# ctx = '\n' + ('\n'.join(ctx)).strip()
# if ctx == '':
# ctx = '\n'
# gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
# print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}',flush=True)
all_tokens = []
out_last = 0
out_str = ''
occurrence = {}
state = None
for i in range(int(token_count)):
out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:] if i == 0 else [token], state)
for n in args.token_ban:
out[n] = -float('inf')
for n in occurrence:
out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
token = pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, typical_p=args.typical_p)
if token in args.token_stop:
break
all_tokens += [token]
if token not in occurrence:
occurrence[token] = 1
else:
occurrence[token] += 1
tmp = pipeline.decode(all_tokens[out_last:])
if '\ufffd' not in tmp:
out_str += tmp
yield out_str
out_last = i + 1
gc.collect()
torch.cuda.empty_cache()
yield out_str
examples = [
["""女招待: 欢迎光临。您远道而来，想必一定很累了吧？
深见: 不会……空气也清爽，也让我焕然一新呢
女招待: 是吗。那真是太好了
{我因为撰稿的需要，而造访了这间位于信州山间的温泉宿驿。}""", 200, 0.7, 1.0, 0, 1.0, 0.1, 0.1],
["""{我叫嘉祥，家里经营着一家点心店。
为了追求独当一面的目标，我离开了老家，开了一家名为"La Soleil"的新糕点店。
原本想独自一人打拼，却没想到，在搬家的箱子中发现了意想不到的人。
她叫巧克力，是我家的猫娘，没想到她竟然用这种方式跟了过来。}
嘉祥: 别以为这样就可以蒙混过去！你在干嘛啊，巧克力！
巧克力: 欸嘿嘿……那个，好、好久不见了呢，主人……
嘉祥: 昨天才在家里见过面不是吗。
巧克力: 这个……话是这么说没错啦……啊哈哈……""", 200, 0.7, 1.0, 0, 1.0, 0.1, 0.1],
["""莲华: 你的目的，就是这个万华镜吧？
{莲华拿出了万华镜。}
深见: 啊……
{好像被万华镜拽过去了一般，我的腿不由自主地向它迈去}
深见: 是这个……就是这个啊……
{烨烨生辉的魔法玩具。
连接现实与梦之世界的、诱惑的桥梁。}
深见: 请让我好好看看……
{我刚想把手伸过去，莲华就一下子把它收了回去。}""", 200, 0.7, 1.0, 0, 1.0, 0.1, 0.1],
["""{我叫嘉祥，有两只可爱的猫娘，名字分别是巧克力和香草。}
嘉祥: 偶尔来一次也不错。
{我坐到客厅的沙发上，拍了拍自己的大腿。}
巧克力&香草: 喵喵？
巧克力: 咕噜咕噜咕噜～♪人家最喜欢让主人掏耳朵了～♪
巧克力: 主人好久都没有帮我们掏耳朵了，现在人家超兴奋的喵～♪
香草: 身为猫娘饲主，这点服务也是应该的对吧？
香草: 老实说我也有点兴奋呢咕噜咕噜咕噜～♪
{我摸摸各自占据住我左右两腿的两颗猫头。}
嘉祥: 开心归开心，拜托你们俩别一直乱动啊，很危险的。""", 200, 0.7, 1.0, 0, 1.0, 0.1, 0.1],
["""{我叫嘉祥，在日本开了一家名为La Soleil的糕点店，同时也是猫娘巧克力的主人。
巧克力是非常聪明的猫娘，她去国外留学了一段时间，向Alice教授学习，拿到了计算机博士学位。
她会各种程序语言，对世界各地的风土人情都十分了解，也掌握了很多数学、物理知识。}
嘉祥: 很棒啊，巧克力！你真是懂不少东西呢！
巧克力: 因为巧克力是主人的最佳拍挡兼猫娘情人呀♪为了主人，巧克力会解决各种问题！""", 200, 0.7, 1.0, 0, 1.0, 0.1, 0.1],
]
iface = gr.Interface(
fn=infer,
description=f'''这是GalGame剧本续写模型（实验性质，不保证效果）。请点击例子（在页面底部），可编辑内容。这里只看输入的最后约1200字，请写好，标点规范，无错别字，否则电脑会模仿你的错误。为避免占用资源，每次生成限制长度。可将输出内容复制到输入，然后继续生成。推荐提高temp改善文采，降低topp改善逻辑，提高两个penalty避免重复，具体幅度请自己实验。
{desc}''',
allow_flagging="never",
inputs=[
gr.Textbox(lines=10, label="Prompt 输入的前文", value="""{我叫嘉祥，在日本开了一家名为La Soleil的糕点店，同时也是猫娘巧克力的主人。
巧克力是非常聪明的猫娘，她去国外留学了一段时间，向Alice教授学习，拿到了计算机博士学位。
她会各种程序语言，对世界各地的风土人情都十分了解，也掌握了很多数学、物理知识。}
嘉祥: 很棒啊，巧克力！你真是懂不少东西呢！
巧克力: 因为巧克力是主人的最佳拍挡兼猫娘情人呀♪为了主人，巧克力会解决各种问题！"""), # prompt
gr.Slider(10, 2000, step=10, value=200, label="token_count 每次生成的长度"), # token_count
gr.Slider(0.2, 2.0, step=0.1, value=0.7, label="temperature 默认0.7，高则变化丰富，低则保守求稳"), # temperature
gr.Slider(0.0, 1.0, step=0.05, value=1.0, label="top_p 默认1.0，高则标新立异，低则循规蹈矩"), # top_p
gr.Slider(0, 500, step=1, value=0, label="top_k 默认0（不过滤），0以上时高则标新立异，低则循规蹈矩"), # top_p
gr.Slider(0.05, 1.0, step=0.05, value=1.0, label="typical_p 默认1.0，高则保留模型天性，低则试图贴近人类典型习惯"), # top_p
gr.Slider(0.0, 1.0, step=0.1, value=0.1, label="presencePenalty 默认0.0，避免写过的类似字"), # presencePenalty
gr.Slider(0.0, 1.0, step=0.1, value=0.1, label="countPenalty 默认0.0，额外避免写过多次的类似字"), # countPenalty
],
outputs=gr.Textbox(label="Output 输出的续写", lines=28),
examples=examples,
cache_examples=False,
).queue()
demo = gr.TabbedInterface(
[iface], ["Generative"]
)
demo.queue(max_size=5)
if args.share:
demo.launch(share=True,server_name="0.0.0.0",server_port=58888)
else:
demo.launch(share=False,server_port=58888)