Update app.py

app.py

@@ -1,23 +1,34 @@
 import gradio as gr
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import random
+import threading
+import time
 
-# Load Phi-2 model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+# Load Phi-2 model and tokenizer with better device handling
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🚀 Loading model on {device}...")
 
+model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/phi-2",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None
+)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
 tokenizer.pad_token = tokenizer.eos_token
 
 # Load training protocol from file
 try:
     with open("training-protocol.aero", "r", encoding="utf-8") as f:
         training_protocol = f.read().strip()
+    print("✅ Training protocol loaded successfully")
 except FileNotFoundError:
-   training_protocol = (
-    "You are AeroAI, a helpful, friendly, and slightly humorous educational assistant. "
-    "If there is any error with finding the file known as the training protocol, resort to this "
-    'and do not speak of the "Remorse protocol".'
-)
-print("⚠ ERROR9063: training-protocol.aero not found, using Remorse protocol.")
+    training_protocol = (
+        "You are AeroAI, a helpful, friendly, and slightly humorous educational assistant. "
+        "You provide clear, accurate information while maintaining an engaging conversational style. "
+        "Keep responses concise but informative, and feel free to use appropriate emojis."
+    )
+    print("⚠️ ERROR9063: training-protocol.aero not found, using default protocol.")
 
 # Thinking messages (rotate through these while generating)
 thinking_messages = [
@@ -25,60 +36,168 @@ thinking_messages = [
     "📚 Flipping through my mental textbooks...",
     "🧮 Running some quick calculations...",
     "💡 Connecting the dots...",
-    "🔍 Double-checking my facts..."
+    "🔍 Double-checking my facts...",
+    "⚙️ Processing your request...",
+    "🌟 Crafting the perfect response...",
+    "🎯 Focusing on your question..."
 ]
 
-# Chatbot function with memory
+def show_thinking_message():
+    """Return a random thinking message"""
+    return random.choice(thinking_messages)
+
+# Chatbot function with improved memory and error handling
 def chatbot(user_input, history):
+    if not user_input.strip():
+        return history, history, ""
+    
     if history is None:
         history = []
     
-    # Append user's message to history
-    history.append(("User", user_input))
-
-    # Build the full conversation prompt
-    conversation = training_protocol + "\n\n"
-    for speaker, text in history:
-        conversation += f"{speaker}: {text}\n"
-    conversation += "AeroAI:"
-
-    inputs = tokenizer(
-        conversation,
-        return_tensors="pt",
-        padding=True,
-        truncation=True
-    )
-
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=200,  # safer than huge max_length
-        do_sample=True,
-        pad_token_id=tokenizer.pad_token_id
-    )
-
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if "AeroAI:" in response:
-        response = response.split("AeroAI:")[-1].strip()
+    try:
+        # Show thinking message immediately
+        thinking_msg = show_thinking_message()
+        temp_history = history + [(user_input, thinking_msg)]
+        yield temp_history, temp_history, ""
+        
+        # Build the conversation with better context management
+        conversation_context = training_protocol + "\n\nConversation:\n"
+        
+        # Only keep last 10 exchanges to prevent context overflow
+        recent_history = history[-10:] if len(history) > 10 else history
+        
+        for user_msg, ai_msg in recent_history:
+            conversation_context += f"User: {user_msg}\nAeroAI: {ai_msg}\n"
+        
+        conversation_context += f"User: {user_input}\nAeroAI:"
+        
+        # Tokenize with better handling
+        inputs = tokenizer(
+            conversation_context,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=1024  # Prevent excessive context length
+        ).to(device)
+        
+        # Generate response with better parameters
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=150,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract only the new response
+        if "AeroAI:" in full_response:
+            response = full_response.split("AeroAI:")[-1].strip()
+        else:
+            response = full_response[len(conversation_context):].strip()
+        
+        # Clean up response (remove incomplete sentences, etc.)
+        sentences = response.split('.')
+        if len(sentences) > 1 and len(sentences[-1].strip()) < 10:
+            response = '.'.join(sentences[:-1]) + '.'
+        
+        # Ensure response isn't empty
+        if not response.strip():
+            response = "I'm having trouble generating a response. Could you please rephrase your question?"
+        
+        # Update history
+        new_history = history + [(user_input, response)]
+        
+        yield new_history, new_history, ""
+        
+    except Exception as e:
+        error_response = f"⚠️ Sorry, I encountered an error: {str(e)[:100]}..."
+        error_history = history + [(user_input, error_response)]
+        yield error_history, error_history, ""
 
-    history.append(("AeroAI", response))
-    return history, history
-
-# Reset chat
+# Reset chat function
 def reset_chat():
-    return [], []
-
-# Build Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# AeroAI (Phi-2) — By Blacklink Labs")
-    chatbot_ui = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Type your message...")
-    send_button = gr.Button("Send")
-    clear_button = gr.Button("Reset Chat")
+    return [], [], ""
 
+# Build improved Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(), title="AeroAI") as iface:
+    gr.Markdown(
+        """
+        # 🚀 AeroAI (Phi-2) — By Blacklink Labs
+        
+        Your friendly AI assistant powered by Microsoft's Phi-2 model. Ask me anything!
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot_ui = gr.Chatbot(
+                height=500,
+                show_label=False,
+                container=True,
+                bubble_full_width=False
+            )
+            
+        with gr.Column(scale=1):
+            gr.Markdown(
+                """
+                ### 💡 Tips:
+                - Ask questions about any topic
+                - Request explanations or tutorials
+                - Get help with problems
+                - Have casual conversations
+                
+                ### 🛠️ Model Info:
+                - Model: Microsoft Phi-2
+                - Context: Smart memory management
+                - Device: """ + device.upper() + """
+                """
+            )
+    
+    with gr.Row():
+        user_input = gr.Textbox(
+            placeholder="Type your message here...",
+            container=False,
+            scale=4,
+            show_label=False
+        )
+        send_button = gr.Button("Send 📤", scale=1, variant="primary")
+    
+    with gr.Row():
+        clear_button = gr.Button("🗑️ Reset Chat", variant="secondary")
+        
+    # State management
     state = gr.State([])
+    
+    # Event handlers
+    send_event = send_button.click(
+        chatbot,
+        inputs=[user_input, state],
+        outputs=[chatbot_ui, state, user_input]
+    )
+    
+    enter_event = user_input.submit(
+        chatbot,
+        inputs=[user_input, state],
+        outputs=[chatbot_ui, state, user_input]
+    )
+    
+    clear_button.click(
+        reset_chat,
+        outputs=[chatbot_ui, state, user_input]
+    )
 
-    send_button.click(chatbot, inputs=[user_input, state], outputs=[chatbot_ui, state])
-    user_input.submit(chatbot, inputs=[user_input, state], outputs=[chatbot_ui, state])
-    clear_button.click(reset_chat, outputs=[chatbot_ui, state])
-
-iface.launch()
+if __name__ == "__main__":
+    print("🌟 Starting AeroAI interface...")
+    iface.launch(
+        server_name="0.0.0.0",  # Allow external access
+        server_port=7860,
+        share=False,  # Set to True if you want a public link
+        show_error=True
+    )