Update app.py

app.py

@@ -1,86 +1,77 @@
 import gradio as gr
 import random
-from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer
-model_name = "facebook/blenderbot-400M-distill"
-tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
-model = BlenderbotForConditionalGeneration.from_pretrained(model_name)
+# Load the faster, lightweight model (DialoGPT-small for speed)
+model_name = "microsoft/DialoGPT-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Persistent user details (Memory)
-user_memory = {
-    "name": "Friend",
-    "age": "unknown",
-    "city": "unknown",
-    "likes": "unknown",
-    "favorite": "unknown"
-}
+# Memory for user details (name, age, city, likes, favorites)
+user_memory = {"name": "Friend", "age": "unknown", "city": "unknown", "likes": "unknown", "favorite": "unknown"}
 
-# Friendly responses
-friendly_replies = [
-    "That sounds awesome! 😃",
-    "You're such a cool person! 😊",
-    "Wow, tell me more! 😎",
-    "Haha, I love chatting with you! 😆",
-    "You always have interesting things to say! 🤩"
-]
-
-# Questions to keep the chat fun
-friendly_questions = [
-    "What's something fun you did today? 🎉",
-    "If you could visit any place, where would it be? 🌍",
-    "What’s your favorite hobby? 🎨⚽",
-    "What’s the best advice you've ever received? 🧠",
-    "If you had a superpower, what would it be? ⚡"
-]
-
-# Chat function
-def chat(user_input, history=[]):
+def chat(user_input, chat_history=[]):
     global user_memory
-
-    # Check for personal detail updates
-    for key in user_memory.keys():
-        if key in user_input.lower():
-            user_memory[key] = user_input.split(key)[-1].strip(" ?.!@,#$")
-            return f"Got it! I'll remember your {key} is {user_memory[key]} 😊"
-
-    # Answer personal detail questions
-    for key, value in user_memory.items():
-        if f"what is your {key}" in user_input.lower():
-            return f"Your {key} is {value}! 😊"
-
-    # Generate AI response
-    inputs = tokenizer(user_input, return_tensors="pt")
-    reply_ids = model.generate(**inputs, max_length=100)
-    response = tokenizer.decode(reply_ids[0], skip_special_tokens=True)
-
-    # Add a friendly response and sometimes ask a fun question
-    response += " " + random.choice(friendly_replies)
-    if random.random() < 0.3:
-        response += " " + random.choice(friendly_questions)
+    user_input_lower = user_input.lower()
     
-    return response
-
-# Gradio UI with enhanced design
-with gr.Blocks(theme=gr.themes.Soft()) as iface:
-    gr.Markdown("""
-    <h1 style='text-align: center; color: #4CAF50;'>🌟 AI Friend Chatbot - Your Virtual Buddy! 🌟</h1>
-    <p style='text-align: center;'>Let's chat and have fun together! 😊</p>
-    """)
+    # Capture user details for memory
+    if "my name is" in user_input_lower:
+        user_memory["name"] = user_input.split("my name is")[-1].strip(" ?.!\n")
+        return f"Nice to meet you, {user_memory['name']}! 😊"
     
-    chatbot = gr.Chatbot(label="Chat with Your AI Friend 🤖", bubble_full_width=False)
+    if "how old am i" in user_input_lower:
+        return f"You told me you are {user_memory['age']} years old! 🎂"
+    if "i am" in user_input_lower and "years old" in user_input_lower:
+        user_memory["age"] = user_input.split("i am")[-1].split("years old")[0].strip()
+        return f"Got it! You are {user_memory['age']} years old. 🎉"
     
-    with gr.Row():
-        user_input = gr.Textbox(placeholder="Say something... ✨", label="You 💬", show_label=False)
-        send_btn = gr.Button("🚀 Send")
-
-    def respond(message, history):
-        response = chat(message)
-        history.append((message, response))
-        return history, ""
+    if "where do i live" in user_input_lower:
+        return f"You told me you live in {user_memory['city']}! 🏡"
+    if "i live in" in user_input_lower:
+        user_memory["city"] = user_input.split("i live in")[-1].strip(" ?.!\n")
+        return f"Awesome! {user_memory['city']} sounds like a great place! 🌍"
+    
+    if "what do i like" in user_input_lower:
+        return f"You said you like {user_memory['likes']}! 😃"
+    if "i like" in user_input_lower:
+        user_memory["likes"] = user_input.split("i like")[-1].strip(" ?.!\n")
+        return f"Nice! {user_memory['likes']} sounds amazing! 😍"
+    
+    if "what is my favorite" in user_input_lower:
+        return f"Your favorite is {user_memory['favorite']}! 🎉"
+    if "my favorite is" in user_input_lower:
+        user_memory["favorite"] = user_input.split("my favorite is")[-1].strip(" ?.!\n")
+        return f"Cool! {user_memory['favorite']} is a great choice! 🏆"
+    
+    # Chatbot response generation (with streaming for speed)
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
+    chat_history_ids = model.generate(input_ids, max_length=50, pad_token_id=tokenizer.eos_token_id)
+    bot_response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    
+    return bot_response
 
-    send_btn.click(respond, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
-    user_input.submit(respond, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+# Gradio UI with attractive design
+def chat_ui():
+    with gr.Blocks(theme=gr.themes.Soft()) as ui:
+        gr.Markdown("""
+        <h1 style='text-align: center; color: blue;'>🤖 AI Friend Chatbot</h1>
+        <p style='text-align: center;'>Your friendly AI companion! Let's chat. 😊</p>
+        """)
+        
+        chatbot = gr.Chatbot(label="Your AI Friend", bubble_full_width=False)
+        user_input = gr.Textbox(placeholder="Type a message...", label="You", show_label=False)
+        send_btn = gr.Button("Send 💬")
+        
+        def respond(message, history):
+            response = chat(message)
+            history.append((message, response))
+            return history, ""
+        
+        send_btn.click(respond, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+        user_input.submit(respond, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+    
+    return ui
 
-# Launch the Chatbot
-iface.launch()
+# Launch the chatbot
+chat_ui().launch()