Update app.py

app.py

@@ -1,40 +1,65 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Better chatbot model
-chatbot = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-small"
-)
+# Load chat model
+model_name = "microsoft/DialoGPT-medium"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Chat function with memory
 def reply(message, history):
     if not message.strip():
         return "Please enter a message."
 
-    prompt = f"User: {message}\nAssistant:"
+    # Build conversation history
+    chat_history_ids = None
+
+    for user, bot in history:
+        user_ids = tokenizer.encode(user + tokenizer.eos_token, return_tensors="pt")
+        bot_ids = tokenizer.encode(bot + tokenizer.eos_token, return_tensors="pt")
+        chat_history_ids = (
+            user_ids if chat_history_ids is None
+            else torch.cat([chat_history_ids, user_ids], dim=-1)
+        )
+        chat_history_ids = torch.cat([chat_history_ids, bot_ids], dim=-1)
+
+    # Add current user message
+    new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
+
+    input_ids = (
+        new_input_ids if chat_history_ids is None
+        else torch.cat([chat_history_ids, new_input_ids], dim=-1)
+    )
 
-    result = chatbot(
-        prompt,
-        max_new_tokens=50,
+    # Generate response
+    output_ids = model.generate(
+        input_ids,
+        max_length=1000,
+        pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
-        temperature=0.7
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7,
+        repetition_penalty=1.2  # 🔥 stops repetition
     )
 
-    output = result[0]["generated_text"]
+    # Extract only new response
+    response_ids = output_ids[:, input_ids.shape[-1]:]
+    response = tokenizer.decode(response_ids[0], skip_special_tokens=True)
 
-    # Extract only assistant response
-    if "Assistant:" in output:
-        response = output.split("Assistant:")[-1].strip()
-    else:
-        response = output.strip()
+    # Clean fallback
+    if response.strip() == "":
+        response = "I'm here! How can I help you?"
 
-    return response  # ✅ ONLY RETURN STRING
+    return response
 
 # Chat UI
 demo = gr.ChatInterface(
     fn=reply,
     title="💬 Smart Dialogue System",
-    description="Chatbot using FLAN-T5"
+    description="Chatbot using DialoGPT (context-aware)"
 )
 
 demo.launch(server_name="0.0.0.0", server_port=7860)