Update app.py

app.py

@@ -1,60 +1,60 @@
 import streamlit as st
 import tempfile
 import wntr
-import os
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# Load model once
 @st.cache_resource
 def load_llm():
-    model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
+    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"  # You can swap this with another DeepSeek model
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True)
     return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
 
-# Load LLM
 llm = load_llm()
 
-st.title("💧 EPANET + LLM (via WNTR + Qwen LLM)")
+st.title("💧 EPANET + WNTR + DeepSeek LLM Assistant")
 
-# Upload .inp file
-uploaded_file = st.file_uploader("Upload EPANET .inp file", type=["inp"])
+uploaded_file = st.file_uploader("Upload your EPANET .inp file", type=["inp"])
 
-# Initialize the WNTR model
 wn = None
 if uploaded_file:
     with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as tmp_file:
         tmp_file.write(uploaded_file.read())
         inp_path = tmp_file.name
     wn = wntr.network.WaterNetworkModel(inp_path)
+    st.success("Water network model loaded successfully.")
+
+question = st.text_input("Ask a question about your water network model")
+
+if st.button("Generate Python Code") and wn and question:
+    prompt = f"""
+You are a Python expert using the WNTR library for water network simulations.
+Given a WNTR water network model `wn`, generate a Python function called `answer()` that answers this question:
 
-# Ask user question
-user_question = st.text_input("Ask a question about your water network model")
+Question: {question}
 
-if st.button("Generate Code") and wn and user_question:
-    prompt = f"""You are a Python expert using the WNTR library. Given a water network model `wn`, answer this question:
-Question: {user_question}
-Generate a Python function called `answer()` that computes and returns the result as a variable named `result`.
-Only include valid Python code. Do not include markdown, explanations, or text outside the function."""
+The function must use the `wn` model, store the final answer in a variable called `result`, and return it.
+Only output valid Python code. Do not include markdown or explanations.
+"""
 
     try:
         response = llm(prompt)[0]["generated_text"]
-        # Optional: Clean hallucinated prompt part
-        if "def answer" in response:
-            response = response.split("def answer", 1)[-1]
-            generated_code = "def answer" + response
+        code_start = response.find("def answer")
+        if code_start != -1:
+            code = response[code_start:]
+            st.subheader("🧠 Generated Code")
+            st.code(code, language="python")
+
+            local_vars = {"wn": wn}
+            try:
+                exec(code, local_vars)
+                result = local_vars["answer"]()
+                st.subheader("📤 Output")
+                st.success(result)
+            except Exception as e:
+                st.subheader("📤 Output")
+                st.error(f"Error executing function: {e}")
         else:
-            raise ValueError("Function definition not found in LLM output.")
-
-        st.markdown("### 🧠 Generated Code")
-        st.code(generated_code, language="python")
-
-        local_env = {"wn": wn}
-        exec(generated_code, local_env)
-        result = local_env["answer"]()
-        st.markdown("### 📤 Output")
-        st.success(result)
+            st.error("Could not extract Python function from LLM response.")
     except Exception as e:
-        st.markdown("### 📤 Output")
-        st.error(f"Error running generated code: {e}")
+        st.error(f"Error querying DeepSeek model: {e}")