import gradio as gr
import os, gc, torch
from datetime import datetime
from huggingface_hub import hf_hub_download
from pynvml import *
nvmlInit()
gpu_h = nvmlDeviceGetHandleByIndex(0)
ctx_limit = 1024
title = "RWKV-4-Pile-14B-20230313-ctx8192-test1050"
desc = f'''Links:<a href='https://github.com/BlinkDL/ChatRWKV' target="_blank" style="margin:0 0.5em">ChatRWKV</a><a href='https://github.com/BlinkDL/RWKV-LM' target="_blank" style="margin:0 0.5em">RWKV-LM</a><a href="https://pypi.org/project/rwkv/" target="_blank" style="margin:0 0.5em">RWKV pip package</a><a href="https://huggingface.co/spaces/BlinkDL/Raven-RWKV-7B" target="_blank" style="margin:0 0.5em">Raven 7B (alpaca-style)</a>
'''

os.environ["RWKV_JIT_ON"] = '1'
os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)

from rwkv.model import RWKV
model_path = hf_hub_download(repo_id="BlinkDL/rwkv-4-pile-14b", filename=f"{title}.pth")
model = RWKV(model=model_path, strategy='cuda fp16i8 *24 -> cuda fp16')

from rwkv.utils import PIPELINE, PIPELINE_ARGS
pipeline = PIPELINE(model, "20B_tokenizer.json")

def infer(
        ctx,
        token_count=10,
        temperature=1.0,
        top_p=0.8,
        presencePenalty = 0.1,
        countPenalty = 0.1,
):
    args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_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(' ')
    if ctx.endswith('\n'):
        ctx = f'\n{ctx.strip()}\n'
    else:
        ctx = f'\n{ctx.strip()}'

    gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
    print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
    
    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)
        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.strip()
            out_last = i + 1
    gc.collect()
    torch.cuda.empty_cache()
    yield out_str.strip()

examples = [
    ["Expert Questions & Helpful Answers\nAsk Research Experts\nQuestion:\nHow can we eliminate poverty?\n\nFull Answer:\n", 150, 1.0, 0.7, 0.2, 0.2],
    ["Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n", 150, 1.0, 0.7, 0.2, 0.2],
    ['''Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Generate a list of adjectives that describe a person as brave.

### Response:
''', 150, 1.0, 0.2, 0.5, 0.5],
    ['''Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
Arrange the given numbers in ascending order.

### Input:
2, 4, 0, 8, 3

### Response:
''', 150, 1.0, 0.2, 0.5, 0.5],
    ["Ask Expert\n\nQuestion:\nWhat are some good plans for world peace?\n\nExpert Full Answer:\n", 150, 1.0, 0.7, 0.2, 0.2],
    ["Q & A\n\nQuestion:\nWhy is the sky blue?\n\nDetailed Expert Answer:\n", 150, 1.0, 0.7, 0.2, 0.2],
    ["Dear sir,\nI would like to express my boundless apologies for the recent nuclear war.", 150, 1.0, 0.7, 0.2, 0.2],
    ["Here is a shell script to find all .hpp files in /home/workspace and delete the 3th row string of these files:", 150, 1.0, 0.7, 0.1, 0.1],
    ["Building a website can be done in 10 simple steps:\n1.", 150, 1.0, 0.7, 0.2, 0.2],
    ["A Chinese phrase is provided: 百闻不如一见。\nThe masterful Chinese translator flawlessly translates the phrase into English:", 150, 1.0, 0.5, 0.2, 0.2],
    ["I believe the meaning of life is", 150, 1.0, 0.7, 0.2, 0.2],
    ["Simply put, the theory of relativity states that", 150, 1.0, 0.5, 0.2, 0.2],
]

chat_intro = '''The following is a coherent verbose detailed conversation between an AI girl named <|bot|> and <|user|>. One day, they meet at a café.
Note the following important facts about <|bot|>:
1. <|bot|> is very intelligent, creative and friendly.
2. <|bot|> likes to tell <|user|> a lot about herself and her opinions.
3. <|bot|> usually gives <|user|> kind, helpful and informative advices.

<|user|>: Hello, how are you doing?

<|bot|>: Hi! Thanks, I'm fine. What about you?

<|user|>: I am fine. It's nice to see you. Look, here is a store selling tea and juice. We can go and take a look. Would you like to chat with me for a while?

<|bot|>: Sure. Let's go inside. What would you like to talk about? I'm listening.
'''

def user(message, chatbot):
    chatbot = chatbot or []
    print(f"User: {message}")
    return "", chatbot + [[message, None]]

def chat(
        prompt,
        user,
        bot,
        chatbot,
        history,
        temperature=1.0,
        top_p=0.8,
        presence_penalty=0.1,
        count_penalty=0.1,
):
    args = PIPELINE_ARGS(temperature=max(0.2, float(temperature)), top_p=float(top_p),
                         alpha_frequency=float(count_penalty),
                         alpha_presence=float(presence_penalty),
                         token_ban=[],  # ban the generation of some tokens
                         token_stop=[])  # stop generation whenever you see any token here

    message = chatbot[-1][0]
    message = message.strip().replace('\r\n','\n').replace('\n\n','\n')
    ctx = f"{user}: {message}\n\n{bot}:"

    # gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
    # print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')

    if not history:
        prompt = prompt.replace("<|user|>", user.strip())
        prompt = prompt.replace("<|bot|>", bot.strip())
        prompt = prompt.strip()
        prompt = f"\n{prompt}\n\n"

        out, state = model.forward(pipeline.encode(prompt), None)
        history = [state, []]
        print("History reloaded.")

    [state, all_tokens] = history
    out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:], state)

    print("Bot: ", end='')

    begin = len(all_tokens)
    out_last = begin
    out_str: str = ''
    occurrence = {}
    for i in range(300):
        if i <= 0:
            nl_bias = -float('inf')
        elif i <= 30:
            nl_bias = (i - 30) * 0.1
        elif i <= 130:
            nl_bias = 0
        else:
            nl_bias = (i - 130) * 0.25
        out[187] += nl_bias
        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)
        next_tokens = [token]
        if token == 0:
            next_tokens = pipeline.encode('\n\n')
        all_tokens += next_tokens

        if token not in occurrence:
            occurrence[token] = 1
        else:
            occurrence[token] += 1

        out, state = model.forward(next_tokens, state)

        tmp = pipeline.decode(all_tokens[out_last:])
        if '\ufffd' not in tmp:
            print(tmp, end='', flush=True)
            out_last = begin + i + 1
            out_str += tmp

            chatbot[-1][1] = out_str.strip()
            history = [state, all_tokens]
            yield chatbot, history

        out_str = pipeline.decode(all_tokens[begin:])
        out_str = out_str.replace("\r\n", '\n').replace('\\n', '\n')

        if '\n\n' in out_str:
            break

    gc.collect()
    torch.cuda.empty_cache()

    chatbot[-1][1] = out_str.strip()
    history = [state, all_tokens]
    yield chatbot, history

with gr.Blocks(title=title) as demo:
    with gr.Tab("Generative"):
        gr.Markdown(f'''{desc} *** <b>Please try examples first (bottom of page)</b> *** (edit them to your own question).\nDemo limited to ctxlen {ctx_limit}.''', label="Description")
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(lines=10, label="Prompt", value="Here's a short cyberpunk sci-fi adventure story. The story's main character is an artificial human created by a company called OpenBot.\n\nThe Story:\n")
                with gr.Row():
                    submit = gr.Button("Submit", variant="primary")
                    clear = gr.Button("Clear", variant="secondary")
                token_count = gr.Slider(10, 200, label="Max Tokens", step=10, value=150)
                temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
                top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.7)
                presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0.2)
                count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0.2)
            with gr.Column():
                output = gr.Textbox(label="Generated Output", lines=32)
        data = gr.Dataset(components=[prompt, token_count, temperature, top_p, presence_penalty, count_penalty], samples=examples, label="Example Prompts", headers=["Prompt", "Max Tokens", "Temperature", "Top P", "Presence Penalty", "Count Penalty"])
        submit.click(infer, [prompt, token_count, temperature, top_p, presence_penalty, count_penalty], [output])
        clear.click(lambda: None, [], [output])
        data.click(lambda x: x, [data], [prompt, token_count, temperature, top_p, presence_penalty, count_penalty])
    
    with gr.Tab("Chat"):
        gr.Markdown(f'''{desc} *** <b>Default Chat Scenario: You (Bob) and Bot (Alice) meet at a café.</b> ***\nIf you want to change the scenario, make sure to use an empty new line to separate different people's words. Also, make sure there is no empty new lines within one person's lines. Changes only take effect after clearing.''', label="Description")
        with gr.Row():
            with gr.Column():
                chatbot = gr.Chatbot()
                state = gr.State()
                message = gr.Textbox(label="Message")
                with gr.Row():
                    send = gr.Button("Send", variant="primary")
                    clear = gr.Button("Clear", variant="secondary")
            with gr.Column():
                with gr.Row():
                    user_name = gr.Textbox(lines=1, max_lines=1, label="User Name", value="Bob")
                    bot_name = gr.Textbox(lines=1, max_lines=1, label="Bot Name", value="Alice")
                prompt = gr.Textbox(lines=10, max_lines=50, label="Scenario", value=chat_intro)
                temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
                top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.7)
                presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0.2)
                count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0.2)
        chat_inputs = [
            prompt,
            user_name,
            bot_name,
            chatbot,
            state,
            temperature,
            top_p,
            presence_penalty,
            count_penalty
        ]
        chat_outputs = [chatbot, state]
        message.submit(user, [message, chatbot], [message, chatbot], queue=False).then(chat, chat_inputs, chat_outputs)
        send.click(user, [message, chatbot], [message, chatbot], queue=False).then(chat, chat_inputs, chat_outputs)
        clear.click(lambda: ([], None, ""), [], [chatbot, state, message], queue=False)

demo.queue(max_size=10)
demo.launch(share=False)