Update app.py

app.py

@@ -1,5 +1,5 @@
 # =========================================
-# Dialogue System using DialoGPT (Gradio)
+# Dialogue System (Gradio FIXED VERSION)
 # =========================================
 
 import gradio as gr
@@ -7,9 +7,9 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # -----------------------------
-# 1. Load Model & Tokenizer
+# 1. Load Model
 # -----------------------------
-MODEL_NAME = "microsoft/DialoGPT-medium"
+MODEL_NAME = "microsoft/DialoGPT-small"  # lighter & faster
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
@@ -17,21 +17,21 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(device)
 
-# -----------------------------
-# 2. Chat Function
-# -----------------------------
 chat_history_ids = None
 
-def chat(user_input, history):
+# -----------------------------
+# 2. Chat Function (FIXED)
+# -----------------------------
+def chat(message, history):
     global chat_history_ids
 
-    if not user_input.strip():
-        return history, ""
+    if history is None:
+        history = []
 
-    # Encode user input
-    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(device)
+    # Encode input
+    new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt").to(device)
 
-    # Append to history
+    # Append history
     if chat_history_ids is not None:
         bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
     else:
@@ -40,7 +40,7 @@ def chat(user_input, history):
     # Generate response
     chat_history_ids = model.generate(
         bot_input_ids,
-        max_length=1000,
+        max_length=500,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
         top_k=50,
@@ -48,13 +48,16 @@ def chat(user_input, history):
         temperature=0.7
     )
 
-    # Decode response
+    # Decode
     response = tokenizer.decode(
         chat_history_ids[:, bot_input_ids.shape[-1]:][0],
         skip_special_tokens=True
     )
 
-    history.append((user_input, response))
+    # ✅ NEW FORMAT (IMPORTANT FIX)
+    history.append({"role": "user", "content": message})
+    history.append({"role": "assistant", "content": response})
+
     return history, ""
 
 # -----------------------------
@@ -66,44 +69,25 @@ def reset_chat():
     return [], ""
 
 # -----------------------------
-# 4. Gradio UI
+# 4. UI (NEW CHAT INTERFACE)
 # -----------------------------
-with gr.Blocks(title="Dialogue System") as app:
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 AI Dialogue System")
 
-    gr.Markdown("## 🤖 AI Dialogue System (Chatbot)")
-    gr.Markdown("Chat with an AI using DialoGPT")
+    chatbot = gr.Chatbot(type="messages")  # IMPORTANT
 
-    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your message...")
 
     with gr.Row():
-        user_input = gr.Textbox(
-            placeholder="Type your message...",
-            show_label=False
-        )
+        send = gr.Button("Send")
+        clear = gr.Button("Clear Chat")
 
-    with gr.Row():
-        send_btn = gr.Button("Send")
-        clear_btn = gr.Button("Clear Chat")
-
-    # Button actions
-    send_btn.click(
-        chat,
-        inputs=[user_input, chatbot],
-        outputs=[chatbot, user_input]
-    )
+    send.click(chat, [msg, chatbot], [chatbot, msg])
+    msg.submit(chat, [msg, chatbot], [chatbot, msg])
 
-    user_input.submit(
-        chat,
-        inputs=[user_input, chatbot],
-        outputs=[chatbot, user_input]
-    )
-
-    clear_btn.click(
-        reset_chat,
-        outputs=[chatbot, user_input]
-    )
+    clear.click(reset_chat, outputs=[chatbot, msg])
 
 # -----------------------------
 # 5. Launch
 # -----------------------------
-app.launch()
+demo.launch()