Update app.py

app.py

@@ -2,71 +2,52 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load a lightweight, conversational model
-model_name = "microsoft/DialoGPT-small"
+# Load a lightweight model
+model_name = "mrm8488/GPT-2-finetuned-on-conversational-data"  # ~500MB
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Store chat history
-chat_history_ids = None
-
-# Inference function
-def chat_with_bot(user_input, history=[]):
-    global chat_history_ids
-    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
-
-    if chat_history_ids is not None:
-        bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
-    else:
-        bot_input_ids = new_input_ids
-
-    chat_history_ids = model.generate(
-        bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id,
-        no_repeat_ngram_size=3, do_sample=True, top_k=100, top_p=0.7, temperature=0.8
-    )
-
-    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    history.append((user_input, response))
-    return history, history
-
-# Gradio Interface
-def reset_history():
-    global chat_history_ids
-    chat_history_ids = None
-    return []
-
-theme_css = """
-body {
-  background-color: #FFF3E0;
-  font-family: Arial, sans-serif;
-}
-
-h1 {
-  color: darkorange;
-  text-align: center;
-  font-weight: bold;
-}
-
-.gradio-container {
-  max-width: 900px;
-  margin: auto;
-  padding: 20px;
-  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2);
-  border-radius: 12px;
-  background-color: white;
-}
-"""
-
-with gr.Blocks(css=theme_css) as demo:
-    gr.Markdown("""<h1>🤖 AI Chatbot Assistant</h1>""")
-
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message here...", label="Your Message")
-    clear = gr.Button("🧹 Clear Chat")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Response generation logic
+def generate_response(message, chat_history):
+    try:
+        chat_history = chat_history or []
+        full_prompt = ""
+        for turn in chat_history:
+            full_prompt += f"User: {turn['content']}\n"
+            full_prompt += f"AI: {turn.get('response', '')}\n"
+        full_prompt += f"User: {message}\nAI:"
+
+        inputs = tokenizer.encode(full_prompt, return_tensors="pt").to(device)
+        outputs = model.generate(inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
+        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract only the latest reply after "AI:"
+        response = decoded.split("AI:")[-1].strip()
+
+        # Add to chat history
+        chat_history.append({"role": "user", "content": message})
+        chat_history.append({"role": "assistant", "content": response})
+        return chat_history
+    except Exception as e:
+        return chat_history + [{"role": "assistant", "content": f"⚠️ Error: {str(e)}"}]
+
+# Gradio interface
+with gr.Blocks(css="body { background-color: #FFF9C4; }") as demo:
+    gr.Markdown("<h1 style='color:brown; font-weight:bold;'>🧠 AI Chatbot</h1>")
+    chatbot = gr.Chatbot(label="Talk to AI", type="messages")
+    msg = gr.Textbox(label="Ask something...", placeholder="Type your message here...")
+
+    clear = gr.Button("Clear Chat")
+    footer = gr.Markdown("<p style='text-align: right; color: brown;'>Designed by Mehak Mazhar</p>")
 
     state = gr.State([])
 
-    msg.submit(chat_with_bot, [msg, state], [chatbot, state])
-    clear.click(fn=reset_history, outputs=[chatbot, state])
+    msg.submit(generate_response, [msg, state], [chatbot, state])
+    clear.click(lambda: ([], []), None, [chatbot, state])
 
-demo.launch()
+# Run
+if __name__ == "__main__":
+    demo.launch()