Update app.py

app.py

@@ -1,19 +1,10 @@
 import streamlit as st
 import wntr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import tempfile
 import os
 import re
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-def answer():
-    return 42  # Returns an integer
-
-Only return code, no explanations. Example:
-```python
-def answer():
-    # Your code
-    return result
-@st.cache_resource
 @st.cache_resource
 def load_llm(model_name="tiiuae/falcon-rw-1b"):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -23,11 +14,10 @@ def load_llm(model_name="tiiuae/falcon-rw-1b"):
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return pipe
 
-def simulate_epanet(inp_path):
-    wn = wntr.network.WaterNetworkModel(inp_path)
-    sim = wntr.sim.WNTRSimulator(wn)
-    results = sim.run_sim()
-    return wn, results
+def extract_code(response):
+    match = re.search(r"def answer\(\):[\s\S]+?return.+", response)
+    return match.group(0) if match else "def answer():\n    return 'Error: No valid function found'"
+
 def generate_python_code(pipe, prompt):
     full_prompt = f"""You are a Python expert using the `wntr` library.
 The user asks: "{prompt}"
@@ -36,74 +26,53 @@ You have the following variables:
 - `wn`: a `WaterNetworkModel` object from the EPANET `.inp` file.
 - `results`: simulation output from `wntr.sim`.
 
-You have these variables:
-- `wn`: the EPANET WaterNetworkModel
-- `results`: the simulation output
-
-Write a Python function named `answer()` that answers the question and returns a single result.
+Write a single function named `answer()` that returns the answer as a Python object (int, float, str, or list).
 
-Begin your function with:
+Only return code, no explanations. Example:
+```python
 def answer():
+    # Your code
+    return result
+```
 """
-   esponse = pipe(full_prompt, max_new_tokens=300, do_sample=True)[0]["generated_text"]
-
-# Extract code block only
-code_block = re.findall(r"```python(.*?)```", response, re.DOTALL)
-if code_block:
-    return code_block[0].strip()
-else:
-    # fallback: look for `def answer():` manually
-    return extract_code(response)
-
-def extract_code(text):
-    lines = text.splitlines()
-    code_lines = []
-    in_func = False
-    for line in lines:
-        if line.strip().startswith("def answer("):
-            in_func = True
-        if in_func:
-            code_lines.append(line)
-        if in_func and line.strip() == "":
-            break
-    return "\n".join(code_lines)
+    response = pipe(full_prompt, max_new_tokens=300, do_sample=True)[0]["generated_text"]
+    code_block = re.findall(r"```python(.*?)```", response, re.DOTALL)
+    if code_block:
+        return code_block[0].strip()
+    else:
+        return extract_code(response)
 
 def run_answer_function(code, wn, results):
     local_vars = {"wn": wn, "results": results}
-    exec(code, {}, local_vars)
-    return local_vars["answer"]()
+    try:
+        exec(code, {}, local_vars)
+        return local_vars["answer"]()
+    except Exception as e:
+        return f"⚠️ Error in generated code: {e}"
 
-# --- Streamlit UI ---
-st.title("💧 EPANET + LLM with Streamlit")
+st.set_page_config(page_title="EPANET + LLM", layout="wide")
+st.title("💧 EPANET + LLM (via WNTR + Falcon LLM)")
 
-uploaded_file = st.file_uploader("Upload your EPANET .inp file", type=["inp"])
-user_prompt = st.text_input("Ask a question (e.g., 'What's the max pressure at node 101?')")
+uploaded_file = st.file_uploader("Upload EPANET .inp file", type=["inp"])
+question = st.text_input("Ask a question about your water network model")
 
-if uploaded_file and user_prompt:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp:
-        temp.write(uploaded_file.read())
-        temp_path = temp.name
+if uploaded_file and question:
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as tmp:
+        tmp.write(uploaded_file.read())
+        tmp_path = tmp.name
 
-    with st.spinner("Simulating EPANET model..."):
-        try:
-            wn, results = simulate_epanet(temp_path)
-        except Exception as e:
-            st.error(f"Failed to simulate EPANET model: {e}")
-            st.stop()
+    wn = wntr.network.WaterNetworkModel(tmp_path)
+    sim = wntr.sim.EpanetSimulator(wn)
+    results = sim.run_sim()
 
-    with st.spinner("Loading model..."):
-        pipe = load_llm()
+    pipe = load_llm()
+    code = generate_python_code(pipe, question)
 
-    with st.spinner("Generating Python code..."):
-        code = generate_python_code(pipe, user_prompt)
-        st.code(code, language="python")
+    st.subheader("🧠 Generated Code")
+    st.code(code, language="python")
 
-    with st.spinner("Running generated code..."):
-        try:
-            output = run_answer_function(code, wn, results)
-            st.success("Answer:")
-            st.write(output)
-        except Exception as e:
-            st.error(f"Error running generated code: {e}")
+    st.subheader("📊 Output")
+    output = run_answer_function(code, wn, results)
+    st.write(output)
 
-    os.remove(temp_path)
+    os.remove(tmp_path)