import gradio as gr
from rwkv.model import RWKV
import os, gc, copy, torch, re
from datetime import datetime
from huggingface_hub import hf_hub_download
os.environ["RWKV_JIT_ON"] = '1'
os.environ["RWKV_CUDA_ON"] = '1'
title = "RWKV-4-World-7B-v1-20230626-ctx4096"

model_path = hf_hub_download(repo_id="BlinkDL/rwkv-4-world", filename=f"{title}.pth")
model = RWKV(model=model_path, strategy='cuda fp16i8 *8 -> cuda fp16')



LENGTH=200

def greet(name):
    return "Hello " + name + "!"

def model_complete(text):
    return model.completion(prompt, max_tokens=LENGTH, temperature=0.0)

demo = gr.Interface(fn=model_complete, inputs="text", outputs="text")
    
demo.launch()