Update app.py

app.py

@@ -1,29 +1,51 @@
 import streamlit as st
 from transformers import pipeline
 
-st.set_page_config(page_title="SUST Math AI", page_icon="🎓")
+# 1. Page Configuration for a high-end UI
+st.set_page_config(page_title="SAAD AI ACADEMY | Math Engine", page_icon="🎓", layout="centered")
+
+st.markdown("""
+    <style>
+    .stApp { background: linear-gradient(to bottom, #0f172a, #000000); color: white; }
+    .stButton>button { width: 100%; border-radius: 20px; background: #3b82f6; color: white; }
+    </style>
+    """, unsafe_allow_safe=True)
 
-# This loads the model directly on Hugging Face's server
 @st.cache_resource
-def load_solver():
-    return pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
+def load_engine():
+    # Fast 360M model optimized for CPU
+    return pipeline("text-generation", model="HuggingFaceTB/SmolLM2-360M-Instruct", device=-1)
+
+pipe = load_engine()
 
-pipe = load_solver()
+st.title("🎓 SAAD AI ACADEMY")
+st.subheader("B.Sc. Mathematics Engine V4")
+st.caption("Optimized for Number Theory, Calculus & Linear Algebra")
 
-st.title("🎓 SUST B.Sc. Math Solver")
-st.markdown("Enter your problem below to get a step-by-step solution.")
+# 2. Professor Mode Instructions
+system_instruction = (
+    "You are a Senior Mathematics Professor. Solve the problem with extreme rigor. "
+    "Use LaTeX for all mathematical notation. If it's a divisibility problem, "
+    "always consider Fermat's Little Theorem or Induction."
+)
 
-user_query = st.chat_input("Ex: Find the derivative of (x^2 + 1)^3")
+user_query = st.chat_input("Enter your SUST Math problem...")
 
 if user_query:
     with st.chat_message("user"):
         st.write(user_query)
 
     with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            messages = [
-                {"role": "system", "content": "You are a professional Mathematics Professor. Solve the problem step-by-step with clear logic."},
-                {"role": "user", "content": user_query},
-            ]
-            result = pipe(messages, max_new_tokens=500, do_sample=False)
-            st.markdown(result[0]['generated_text'][-1]['content'])
+        with st.spinner("⏳ Analyzing mathematical structures..."):
+            prompt = f"<|im_start|>system\n{system_instruction}<|im_end|>\n<|im_start|>user\n{user_query}<|im_end|>\n<|im_start|>assistant\n"
+            
+            # Increased tokens for longer proofs
+            result = pipe(prompt, max_new_tokens=800, temperature=0.1, do_sample=False)
+            response = result[0]['generated_text'].split("<|im_start|>assistant\n")[-1]
+            
+            st.markdown(response)
+            
+            # 3. Add a simple Copy/Download button
+            st.download_button("Download Solution", response, file_name="solution.txt")
+
+st.sidebar.info("Tip: Use standard notation like x^2, integrals, or matrices.")