Update tabbed.py

tabbed.py

@@ -60,6 +60,34 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repeat_
         yield history, history
 
 
+def rp_chat(history, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty):
+    history = history or []
+
+    messages = "<|system|>" + system_message + \
+               "\n".join(["\n".join(["<|user|>"+item[0], "<|model|>"+item[1]])
+                        for item in history])
+
+    # remove last space from assistant, some models output a ZWSP if you leave a space
+    messages = messages[:-1]
+
+    history[-1][1] = ""
+    for output in llm(
+            messages,
+            echo=False,
+            stream=True,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repeat_penalty=repeat_penalty,
+            **config['chat']
+    ):
+        answer = output['choices'][0]['text']
+        history[-1][1] += answer
+        # stream the response
+        yield history, history
+
+
 def clear_chat(chat_history_state, chat_message):
     chat_history_state = []
     chat_message = ''
@@ -114,6 +142,7 @@ with gr.Blocks() as demo:
             )
         with gr.Row():
             submit = gr.Button(value="Send message", variant="secondary").style(full_width=True)
+            roleplay = gr.Button(value="Roleplay", variant="secondary").style(full_width=True)
             clear = gr.Button(value="New topic", variant="secondary").style(full_width=False)
             stop = gr.Button(value="Stop", variant="secondary").style(full_width=False)
         with gr.Row():
@@ -136,11 +165,16 @@ with gr.Blocks() as demo:
         ).then(
             fn=chat, inputs=[chat_history_state, system_msg, max_tokens, temperature, top_p, top_k, repeat_penalty], outputs=[chatbot, chat_history_state], queue=True
         )
+        roleplay_click_event = roleplay.click(
+            fn=user, inputs=[message, chat_history_state], outputs=[message, chat_history_state], queue=True
+        ).then(
+            fn=rp_chat, inputs=[chat_history_state, system_msg, max_tokens, temperature, top_p, top_k, repeat_penalty], outputs=[chatbot, chat_history_state], queue=True
+        )
         message_submit_event = message.submit(
             fn=user, inputs=[message, chat_history_state], outputs=[message, chat_history_state], queue=True
         ).then(
             fn=chat, inputs=[chat_history_state, system_msg, max_tokens, temperature, top_p, top_k, repeat_penalty], outputs=[chatbot, chat_history_state], queue=True
         )
-        stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_click_event, message_submit_event], queue=False)
+        stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_click_event, message_submit_event, roleplay_click_event], queue=False)
 
 demo.queue(**config["queue"]).launch(debug=True, server_name="0.0.0.0", server_port=7860)