Update app.py

app.py

@@ -64,6 +64,12 @@ def predict(message, history, system_prompt, temperature, max_new_tokens, top_k,
         yield "".join(outputs)
 
 
+def handle_retry(history, retry_data: gr.RetryData):
+    new_history = history[:retry_data.index]
+    previous_prompt = history[retry_data.index]['content']
+    yield from respond(previous_prompt, new_history)
+
+
 # Load model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
@@ -121,5 +127,6 @@ gr.ChatInterface(
     chatbot=gr.Chatbot(
         scale=1,
         show_copy_button=True
-    )
+    ),
+    chatbot.retry(handle_retry, chatbot, [chatbot])
 ).queue().launch()