1. Fix clearing

2. Add custom chat scenarios.

app.py

@@ -1,9 +1,5 @@
-from rwkv.utils import PIPELINE, PIPELINE_ARGS
-from rwkv.model import RWKV
 import gradio as gr
-import os
-import gc
-import torch
+import os, gc, torch
 from datetime import datetime
 from huggingface_hub import hf_hub_download
 from pynvml import *
@@ -11,35 +7,32 @@ 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>
+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'
-# if '1' then use CUDA kernel for seq mode (much faster)
-os.environ["RWKV_CUDA_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 *20 -> cuda fp16')
-pipeline = PIPELINE(model, "20B_tokenizer.json")
+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,
-        presence_enalty=0.1,
-        count_penalty=0.1,
+        presencePenalty = 0.1,
+        countPenalty = 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_enalty),
-                         token_ban=[0],  # ban the generation of some tokens
-                         token_stop=[])  # stop generation whenever you see any token here
+    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'):
@@ -49,7 +42,7 @@ def infer(
 
     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 = ''
@@ -70,7 +63,7 @@ def infer(
             occurrence[token] = 1
         else:
             occurrence[token] += 1
-
+        
         tmp = pipeline.decode(all_tokens[out_last:])
         if '\ufffd' not in tmp:
             out_str += tmp
@@ -110,92 +103,70 @@ Arrange the given numbers in ascending order.
     ["Simply put, the theory of relativity states that", 150, 1.0, 0.5, 0.2, 0.2],
 ]
 
-# infer_interface = 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()
-
-########################################################################################################
-
-user = "Bob"
-bot = "Alice"
-interface = ":"
-
-chat_intro = f'''
-The following is a coherent verbose detailed conversation between a girl named {bot} and her friend {user}. \
-{bot} is very intelligent, creative and friendly. \
-She is unlikely to disagree with {user}, and she doesn't like to ask {user} questions. \
-She also likes to tell {user} a lot about herself and her opinions, and she usually gives {user} kind, helpful and informative advices.
-
-{user}{interface} Hello, how are you doing?
+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.
 
-{bot}{interface} Hi {user}! Thanks, I'm fine. What about you?
+<|user|>: Hello, how are you doing?
 
-{user}{interface} I am fine. It's nice to see you. Look, here is a store selling tea and juice.
+<|bot|>: Hi! Thanks, I'm fine. What about you?
 
-{bot}{interface} Sure. Let's go inside. I would like to have some Mocha latte, which is my favourite!
-
-{user}{interface} What is it?
-
-{bot}{interface} Mocha latte is usually made with espresso, milk, chocolate, and frothed milk. Its flavors are frequently sweet.
-
-{user}{interface} Sounds tasty. I'll try it next time. Would you like to chat with me for a while?
-
-{bot}{interface} Of course! I'm glad to answer your questions or give helpful advices. You know, I am confident with my expertise. So please go ahead!
+<|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.
 '''
 
-_, intro_state = model.forward(pipeline.encode(chat_intro), None)
-
-def user(user_message, chatbot):
+def user(message, chatbot):
     chatbot = chatbot or []
-    return "", chatbot + [[user_message, None]]
+    print(f"User: {message}")
+    return "", chatbot + [[message, None]]
 
 def chat(
+        prompt,
+        user,
+        bot,
         chatbot,
         history,
-        token_count=10,
         temperature=1.0,
         top_p=0.8,
-        presence_enalty=0.1,
+        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_enalty),
+                         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(' ')
-    message = message.replace('\n', '')
-    ctx = f"{user}{interface} {message}\n\n{bot}{interface}"
+    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}')
+    # gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+    # print(f'vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
 
-    history = history or [intro_state, []]  # [chat, state, all_tokens]
+    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(int(token_count)):
+    for i in range(300):
         if i <= 0:
             nl_bias = -float('inf')
         elif i <= 30:
@@ -239,77 +210,60 @@ def chat(
     history = [state, all_tokens]
     return chatbot, history
 
-# chat_interface = gr.Interface(
-#     fn=chat,
-#     description=f'''You are {user}, bot is {bot}.''',
-#     allow_flagging="never",
-#     inputs = [
-#         gr.Textbox(label="Message"),
-#         "state",
-#         gr.Slider(10, 1000, step=10, value=250),    # 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.8),  # top_p
-#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # presence_penalty
-#         gr.Slider(0.0, 1.0, step=0.1, value=0.2),   # count_penalty
-#     ],
-#     outputs=[
-#         gr.Chatbot(label="Chat Log", color_map=("blue", "pink")),
-#         "state"
-#     ]
-# ).queue()
-
-########################################################################################################
-
-# demo = gr.TabbedInterface(
-#     [infer_interface, chat_interface], ["Generative", "Chat"],
-#     title=title,
-# )
-
-# demo.queue(max_size=10)
-# demo.launch(share=True)
-
 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 use your question). Demo limited to ctxlen {ctx_limit}.''', label="Description")
+        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")
-                token_count = gr.Slider(10, 1000, label="Max Token", step=10, value=250)
+                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.8)
+                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")
-                    clear = gr.Button("Clear")
-                output = gr.Textbox(label="Generated Output", lines=28)
+                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}Scenario: You are Bob, bot is Alice. You meet at a café.''', label="Description")
+        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")
-                    clear = gr.Button("Clear")
+                    send = gr.Button("Send", variant="primary")
+                    clear = gr.Button("Clear", variant="secondary")
             with gr.Column():
-                token_count = gr.Slider(10, 1000, label="Max Token", step=10, value=250)
+                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.8)
+                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)
-        message.submit(user, [message, chatbot], [message, chatbot], queue=False).then(
-            chat, [chatbot, state, token_count, temperature, top_p, presence_penalty, count_penalty], [chatbot, state]
-        )
-        send.click(user, [message, chatbot], [message, chatbot], queue=False).then(
-            chat, [chatbot, state, token_count, temperature, top_p, presence_penalty, count_penalty], [chatbot, state]
-        )
+        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])
 
 demo.queue(max_size=10)