Update app.py

app.py

@@ -59,7 +59,7 @@ def generate_response(prompt, chat_history, max_new_tokens, temperature):
     response = ""
     for token_id in generated_ids.sequences[0][len(model_inputs.input_ids[0]):]:
         response += tokenizer.decode([token_id], skip_special_tokens=True)
-        yield response
+        yield chat_history + [{"role": "assistant", "content": response}]
 
 # Clear Chat History
 # ==================
@@ -67,7 +67,7 @@ def clear_chat():
     """
     Clear the chat history.
     """
-    return []
+    return [], ""
 
 # Gradio Interface
 # =================
@@ -93,12 +93,14 @@ def gradio_interface():
             chat_history.append({"role": "user", "content": message})
             response = ""
             for chunk in generate_response(message, chat_history, max_new_tokens, temperature):
-                response = chunk
-                yield chat_history + [{"role": "assistant", "content": response}]
+                response = chunk[-1]["content"]
+                yield chat_history, ""
             chat_history.append({"role": "assistant", "content": response})
+            yield chat_history, ""
 
-        submit.click(respond, [msg, chatbot, max_new_tokens, temperature], [chatbot])
-        clear.click(clear_chat, None, [chatbot])
+        submit.click(respond, [msg, chatbot, max_new_tokens, temperature], [chatbot, msg])
+        msg.submit(respond, [msg, chatbot, max_new_tokens, temperature], [chatbot, msg])
+        clear.click(clear_chat, None, [chatbot, msg])
 
     demo.launch()