Add chat mode.

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import os, gc, torch
+import os, gc, copy, torch
 from datetime import datetime
 from huggingface_hub import hf_hub_download
 from pynvml import *
@@ -98,27 +98,208 @@ examples = [
     ["You have $100, and your goal is to turn that into as much money as possible with AI and Machine Learning. Please respond with detailed plan.", "", 150, 1.0, 0.5, 0.4, 0.4],
 ]
 
-g = gr.Interface(
-    fn=evaluate,
-    inputs=[
-        gr.components.Textbox(lines=2, label="Instruction", value="Tell me about ravens."),
-        gr.components.Textbox(lines=2, label="Input", placeholder="none"),
-        gr.components.Slider(minimum=10, maximum=200, step=10, value=150), # token_count
-        gr.components.Slider(minimum=0.2, maximum=2.0, step=0.1, value=1.0), # temperature
-        gr.components.Slider(minimum=0, maximum=1, step=0.05, value=0.5), # top_p
-        gr.components.Slider(0.0, 1.0, step=0.1, value=0.4),  # presencePenalty
-        gr.components.Slider(0.0, 1.0, step=0.1, value=0.4),  # countPenalty        
-    ],
-    outputs=[
-        gr.inputs.Textbox(
-            lines=5,
-            label="Output",
-        )
-    ],
-    title=f"🐦Raven - {title}",
-    description="Raven is [RWKV 7B](https://github.com/BlinkDL/ChatRWKV) 100% RNN [RWKV-LM](https://github.com/BlinkDL/RWKV-LM) finetuned to follow instructions. *** Please try examples first (bottom of page) *** (edit them to use your question). Demo limited to ctxlen 1024. It is finetuned on [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca), codealpaca and more. For best results, *** keep you prompt short and clear ***.",
-    examples=examples,
-    cache_examples=False,
-)
-g.queue(concurrency_count=1, max_size=10)
-g.launch(share=False)
+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)
+
+    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:
+    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>🐦Raven - {title}</h1>\n</div>")
+    with gr.Tab("Instruct"):
+        gr.Markdown(f"Raven is [RWKV 7B](https://github.com/BlinkDL/ChatRWKV) 100% RNN [RWKV-LM](https://github.com/BlinkDL/RWKV-LM) finetuned to follow instructions. *** Please try examples first (bottom of page) *** (edit them to use your question). Demo limited to ctxlen 1024. It is finetuned on [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca), codealpaca and more. For best results, *** keep you prompt short and clear ***.")
+        with gr.Row():
+            with gr.Column():
+                instruction = gr.Textbox(lines=2, label="Instruction", value="Tell me about ravens.")
+                input = gr.Textbox(lines=2, label="Input", placeholder="none")
+                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():
+                with gr.Row():
+                    submit = gr.Button("Submit", variant="primary")
+                    clear = gr.Button("Clear", variant="secondary")
+                output = gr.Textbox(label="Output", lines=5)
+        data = gr.Dataset(components=[instruction, input, token_count, temperature, top_p, presence_penalty, count_penalty], samples=examples, label="Example Instructions", headers=["Instruction", "Input", "Max Tokens", "Temperature", "Top P", "Presence Penalty", "Count Penalty"])
+        submit.click(evaluate, [instruction, input, token_count, temperature, top_p, presence_penalty, count_penalty], [output])
+        clear.click(lambda: None, [], [output])
+        data.click(lambda x: x, [data], [instruction, input, token_count, temperature, top_p, presence_penalty, count_penalty])
+    
+    with gr.Tab("Chat"):
+        gr.Markdown(f'''*** <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=True)