Update app.py

app.py

@@ -1,11 +1,12 @@
+
 import gradio as gr
 import gc, copy, re
 import urllib.request
 from rwkv.model import RWKV
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
 
-ctx_limit = 2048
-title = "RWKV-5-World-0.4B-v2-20231113-ctx4096.pth"
+ctx_limit = 4096
+title = "RWKV-5-World-0.1B-v1-20230803-ctx4096.pth"
 
 url = f"https://huggingface.co/BlinkDL/rwkv-5-world/resolve/main/{title}"
 urllib.request.urlretrieve(url, title)
@@ -13,21 +14,22 @@ urllib.request.urlretrieve(url, title)
 model = RWKV(model=title, strategy='cpu bf16')
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
-def generate_prompt(instruction, input=None):
+def generate_prompt(instruction, input=None, history=None):
+    # Parse the chat history into a string of user and assistant messages
+    history_str = ""
+    for pair in history:
+        history_str += f"Instruction: {pair[0]}\n\nAssistant: {pair[1]}\n\n"
+
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n').replace('\n\n','\n')
     input = input.strip().replace('\r\n','\n').replace('\n\n','\n').replace('\n\n','\n')
     if input and len(input) > 0:
-        return f"""Instruction: {instruction}
+        return f"""{history_str}Instruction: {instruction}
 
 Input: {input}
 
 Response:"""
     else:
-        return f"""User: hi
-
-Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.
-
-User: {instruction}
+        return f"""{history_str}User: {instruction}
 
 Assistant:"""
 
@@ -49,6 +51,7 @@ def evaluate(
     top_p=0.5,
     presencePenalty = 0.5,
     countPenalty = 0.5,
+    history=None # Add the history parameter to the evaluate function
 ):
     args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p),
                      alpha_frequency = countPenalty,
@@ -58,7 +61,7 @@ def evaluate(
 
     instruction = re.sub(r'\n{2,}', '\n', instruction).strip().replace('\r\n','\n')
     input = re.sub(r'\n{2,}', '\n', input).strip().replace('\r\n','\n')
-    ctx = generate_prompt(instruction, input)
+    ctx = generate_prompt(instruction, input, history) # Pass the history to the generate_prompt function
     print(ctx + "\n")
     
     all_tokens = []
@@ -111,26 +114,28 @@ def alternative(chatbot, history):
 
 with gr.Blocks(title=title) as demo:
     gr.HTML(f"<div style=\"text-align: center;\">\n<h1>🌍World - {title}</h1>\n</div>")
+
     with gr.Tab("Instruct mode"):
-        gr.Markdown(f"100% RNN RWKV-LM **trained on 100+ world languages**. Demo limited to ctxlen {ctx_limit}. Finetuned on alpaca, gpt4all, codealpaca and more. For best results, <b>keep your prompt short and clear</b>.")
+        gr.Markdown(f"100% RNN RWKV-LM **trained on 100+ natural languages**. Demo limited to ctxlen {ctx_limit}. For best results, <b>keep your prompt short and clear</b>.")
         with gr.Row():
             with gr.Column():
-                instruction = gr.Textbox(lines=2, label="Instruction", value='東京で訪れるべき素晴らしい場所とその紹介をいくつか挙げてください。')
+                instruction = gr.Textbox(lines=2, label="Instruction", value="Please show me a table with a cheat sheet of Python's syntax.")
                 input = gr.Textbox(lines=2, label="Input", placeholder="")
                 token_count = gr.Slider(10, 300, label="Max Tokens", step=10, value=300)
                 temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.2)
                 top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0.3)
                 presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0)
                 count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0.7)
-            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 mode"):
+        chatbot = gr.ChatInterface(fn=evaluate, 
+                                   additional_inputs=[instruction, input, token_count, temperature, top_p, presence_penalty, count_penalty],
+                                   additional_inputs_accordion="Parameters",
+                                   examples=["Hello", "Write a poem about love", "Generate a list of prime numbers"],
+                                   title="RWKV Chatbot",
+                                   description="A chatbot that can generate creative and informative content based on instructions and inputs")
+
 demo.queue(max_size=10)
-demo.launch(share=False)
+demo.launch(share=False)