Add alternative and derail protection.

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 *
@@ -123,6 +123,15 @@ def user(message, chatbot):
     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,
@@ -139,6 +148,9 @@ def chat(
                          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')
@@ -154,11 +166,14 @@ def chat(
         prompt = f"\n{prompt}\n\n"
 
         out, state = model.forward(pipeline.encode(prompt), None)
-        history = [state, []]
+        history = [state, None, []]  # [state, state_pre, tokens]
         print("History reloaded.")
 
-    [state, all_tokens] = history
+    [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='')
 
@@ -208,11 +223,27 @@ def chat(
         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, all_tokens]
+    history = [state, state_pre_0, all_tokens]
     yield chatbot, history
 
 with gr.Blocks(title=title) as demo:
@@ -245,6 +276,7 @@ with gr.Blocks(title=title) as demo:
                 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():
@@ -269,6 +301,7 @@ with gr.Blocks(title=title) as demo:
         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)