Update app.py

app.py

@@ -5,7 +5,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
 tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
 
-# Set pad_token as eos_token
 tokenizer.pad_token = tokenizer.eos_token
 
 # Load training protocol from file
@@ -16,39 +15,66 @@ except FileNotFoundError:
     training_protocol = "You are AeroAI, a helpful, friendly, and slightly humorous educational assistant."
     print("⚠ training-protocol.aero not found, using default protocol.")
 
-# Chatbot function
-def chatbot(input_text):
-    # Combine protocol with user input
-    prompt = f"{training_protocol}\n\nUser: {input_text}\nAeroAI:"
+# Thinking messages (rotate through these while generating)
+thinking_messages = [
+    "🤔 Thinking deeply about your question...",
+    "📚 Flipping through my mental textbooks...",
+    "🧮 Running some quick calculations...",
+    "💡 Connecting the dots...",
+    "🔍 Double-checking my facts..."
+]
+
+# Chatbot function with memory
+def chatbot(user_input, 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(
-        prompt,
+        conversation,
         return_tensors="pt",
         padding=True,
         truncation=True
     )
+
     outputs = model.generate(
         **inputs,
-        max_length=2000000000000,
+        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)
-    
-    # Optional: cut out the protocol part from the output
     if "AeroAI:" in response:
         response = response.split("AeroAI:")[-1].strip()
-    
-    return response
-
-# Gradio Interface
-iface = gr.Interface(
-    fn=chatbot,
-    inputs="text",
-    outputs="text",
-    title="AeroAI (Phi-2)",
-    description="By Blacklink Education in collaboration with Microsoft."
-)
-
-# Launch the interface
+
+    history.append(("AeroAI", response))
+    return history, history
+
+# Reset chat
+def reset_chat():
+    return [], []
+
+# Build Gradio interface
+with gr.Blocks() as iface:
+    gr.Markdown("# AeroAI (Phi-2) — By Blacklink Education")
+    chatbot_ui = gr.Chatbot()
+    user_input = gr.Textbox(placeholder="Type your message...")
+    send_button = gr.Button("Send")
+    clear_button = gr.Button("Reset Chat")
+
+    state = gr.State([])
+
+    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()