Update app.py

app.py

@@ -1,91 +1,51 @@
 import streamlit as st
-import wntr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import tempfile
 import os
+import wntr
+from transformers import pipeline
 
-@st.cache_resource
-def load_llm(model_name="mistralai/Mistral-7B-Instruct-v0.2"):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, torch_dtype="auto", device_map="auto"
-    )
-    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 generate_python_code(pipe, prompt):
-    full_prompt = f"""You are a Python expert using the `wntr` library.
-Here is the user's question: "{prompt}"
-
-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.
-
-Begin your function with:
-def answer():
-"""
-    response = pipe(full_prompt, max_new_tokens=300, do_sample=True)[0]["generated_text"]
-    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)
-
-def run_answer_function(code, wn, results):
-    local_vars = {"wn": wn, "results": results}
-    exec(code, {}, local_vars)
-    return local_vars["answer"]()
-
-# --- Streamlit UI ---
-st.title("💧 EPANET + LLM with Streamlit")
-
-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?')")
-
-if uploaded_file and user_prompt:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp:
-        temp.write(uploaded_file.read())
-        temp_path = temp.name
+# Load LLM
+llm = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", device_map="auto", max_new_tokens=200)
 
-    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()
+st.title("💧 EPANET + LLM (via WNTR + Falcon LLM)")
 
-    with st.spinner("Loading model..."):
-        pipe = load_llm()
+# Upload .inp file
+uploaded_file = st.file_uploader("Upload EPANET .inp file", type=["inp"])
 
-    with st.spinner("Generating Python code..."):
-        code = generate_python_code(pipe, user_prompt)
-        st.code(code, language="python")
+# Initialize the model globally
+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)
 
-    with st.spinner("Running generated code..."):
+# Ask question
+user_question = st.text_input("Ask a question about your water network model")
+
+if st.button("Generate Code") and wn and user_question:
+    # Prompt to LLM
+    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."""
+    
+    llm_response = llm(prompt)[0]["generated_text"]
+    code_start = llm_response.find("def answer")
+    if code_start != -1:
+        generated_code = llm_response[code_start:]
+        st.markdown("### 🧠 Generated Code")
+        st.code(generated_code, language="python")
+
+        # Execute the function in a controlled environment
+        local_env = {"wn": wn}
         try:
-            output = run_answer_function(code, wn, results)
-            st.success("Answer:")
-            st.write(output)
+            exec(generated_code, local_env)
+            result = local_env["answer"]()
+            st.markdown("### 📤 Output")
+            st.success(result)
         except Exception as e:
-            st.error(f"Error running generated code: {e}")
-
-    os.remove(temp_path)
-
-
+            st.markdown("### 📤 Output")
+            st.error(f"Error in generated code: {e}")
+    else:
+        st.error("Failed to extract Python function from LLM response.")