Implement basic chat mode

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import os, gc, torch
+import os, copy, gc, torch
 from datetime import datetime
 from huggingface_hub import hf_hub_download
 from pynvml import *
@@ -7,19 +7,16 @@ 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>
+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"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
+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")
 
@@ -106,28 +103,206 @@ Arrange the given numbers in ascending order.
     ["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 alternative(chatbot, history):
+    if not chatbot or not history:
+        return chatbot, history
+    
+    chatbot[-1][1] = None
+    history[0] = copy.deepcopy(history[1])
+
+    return chatbot, history
+
+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
+    
+    if not chatbot:
+        return chatbot, history
+
+    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, None, []]  # [state, state_pre, tokens]
+        print("History reloaded.")
+
+    [state, _, all_tokens] = history
+    state_pre_0 = copy.deepcopy(state)
 
-iface = gr.Interface(
-    fn=infer,
-    description=f'''{desc} *** <b>Please try examples first (bottom of page)</b> *** (edit them to use your question). Demo limited to ctxlen {ctx_limit}.''',
-    allow_flagging="never",
-    inputs=[
-        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"),  # prompt
-        gr.Slider(10, 200, step=10, value=150),  # token_count
-        gr.Slider(0.2, 2.0, step=0.1, value=1.0),  # temperature
-        gr.Slider(0.0, 1.0, step=0.05, value=0.7),  # top_p
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # presencePenalty
-        gr.Slider(0.0, 1.0, step=0.1, value=0.2),  # countPenalty
-    ],
-    outputs=gr.Textbox(label="Generated Output", lines=28),
-    examples=examples,
-    cache_examples=False,
-).queue()
-
-demo = gr.TabbedInterface(
-    [iface], ["Generative"],
-    title=title,
-)
+    out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:], state)
+    state_pre_1 = copy.deepcopy(state)    # For recovery
+
+    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
+
+        # State recovery
+        if f'{user}:' in out_str or f'{bot}:' in out_str:
+            idx_user = out_str.find(f'{user}:')
+            idx_user = len(out_str) if idx_user == -1 else idx_user
+            idx_bot = out_str.find(f'{bot}:')
+            idx_bot = len(out_str) if idx_bot == -1 else idx_bot
+            idx = min(idx_user, idx_bot)
+
+            if idx < len(out_str):
+                out_str = f" {out_str[:idx].strip()}\n\n"
+                tokens = pipeline.encode(out_str)
+
+                all_tokens = all_tokens[:begin] + tokens
+                out, state = model.forward(tokens, state_pre_1)
+                break
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    chatbot[-1][1] = out_str.strip()
+    history = [state, state_pre_0, 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")
+                    alt = gr.Button("Alternative", variant="secondary")
+                    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)
+        alt.click(alternative, [chatbot, state], [chatbot, state], 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)