Update app.py

app.py

@@ -1,18 +1,18 @@
 # streamlit_app.py
-# app.py
+# streamlit_app.py
 import streamlit as st
 import re
 from sympy import symbols, integrate, exp, pi
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-st.set_page_config(page_title="AI Physics Solver", page_icon="🧠")
+st.set_page_config(page_title="AI Problem Solver By Mathematically Modelling", page_icon="🧠")
 
 x, t = symbols("x t")
 
 def extract_integral(problem_text):
     match = re.search(r'(\d+)\*?[tx]\^(\d+)', problem_text)
-    limits = re.findall(r'[tx]\s*=\s*([\d\.\\\w]+)', problem_text)
+    limits = re.findall(r'[tx]\s*=\s*([\d\.\w]+)', problem_text)
     exp_match = re.search(r'(\d+)e\^([\-\+]?\d+\.?\d*)[tx]', problem_text)
 
     if 'radioactive' in problem_text or 'half-life' in problem_text:
@@ -35,61 +35,61 @@ def extract_integral(problem_text):
     return "Could not parse the integral format."
 
 @st.cache_resource
-def load_deepseek():
-    model_name = "deepseek-ai/deepseek-math-7b-base"
+def load_model():
+    # Change this if you want to fallback to a smaller model on CPU
+    use_light_model = not torch.cuda.is_available()
+
+    model_name = (
+        "deepseek-ai/deepseek-math-7b-base" if not use_light_model
+        else "tiiuae/falcon-7b-instruct"
+    )
+
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    if torch.cuda.is_available():
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-    else:
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None
+    )
 
     return tokenizer, model
 
-def run_deepseek(user_question):
-    tokenizer, model = load_deepseek()
-    solution_steps = """
-### Solution:
-1. Understand the problem and extract known quantities.
-2. Apply relevant physical laws or mathematical formulas.
-3. Solve algebraically or numerically as required.
-4. Clearly present the final answer.
+def run_llm_reasoning(user_question):
+    tokenizer, model = load_model()
+
+    prompt = f"""
+Q: Solve the following physics problem using rigorous mathematical reasoning. Do not skip any steps.
+
+Problem: {user_question}
 
 ### Final Answer Format:
 Final Answer: [VARIABLE] = [ANSWER] [UNIT]
-"""
-    prompt = f"Q: Solve the following physics problem using rigorous mathematical reasoning. Do not skip any steps.\n\nProblem: {user_question}\n\n{solution_steps}\nA:"
-    inputs = tokenizer(prompt, return_tensors="pt")
+A:"""
 
-    # Move inputs to GPU if available
-    if torch.cuda.is_available():
-        inputs = inputs.to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=500,
-            temperature=0.1,
+            temperature=0.2,
             repetition_penalty=1.0,
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.eos_token_id
         )
+
     return tokenizer.decode(outputs[0], skip_special_tokens=True).split("A:")[-1].strip()
 
-# ---------------- UI Layout ----------------
-st.title("🧠 AI Science Solver")
+# ---------------- UI ----------------
+st.title("🧠 AI Physics & Math Solver")
 
-task_type = st.selectbox("Choose Task Type", ["LLM Reasoning (DeepSeek)", "Symbolic Integration"])
+task_type = st.selectbox("Choose Task Type", ["LLM Reasoning (DeepSeek/Fallback)", "Symbolic Integration"])
 user_question = st.text_area("Enter your physics or math question below:")
 
 if st.button("Solve"):
     with st.spinner("Solving..."):
-        if task_type == "LLM Reasoning (DeepSeek)":
-            result = run_deepseek(user_question)
+        if task_type == "LLM Reasoning (DeepSeek/Fallback)":
+            result = run_llm_reasoning(user_question)
         else:
             result = extract_integral(user_question)