Update app.py

app.py

@@ -1,78 +1,64 @@
+
 import streamlit as st
 import wntr
 import tempfile
-import os
-import re
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 @st.cache_resource
-def load_llm(model_name="tiiuae/falcon-rw-1b"):
+def load_llm():
+    model_name = "tiiuae/falcon-7b-instruct"
     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 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}"
-
-You have the following variables:
-- `wn`: a `WaterNetworkModel` object from the EPANET `.inp` file.
-- `results`: simulation output from `wntr.sim`.
-
-Write a single function named `answer()` that returns the answer as a Python object (int, float, str, or list).
-
-Only return code, no explanations. Example:
-```python
-def answer():
-    # Your code
-    return result
-```
-"""
-    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}
-    try:
-        exec(code, {}, local_vars)
-        return local_vars["answer"]()
-    except Exception as e:
-        return f"⚠️ Error in generated code: {e}"
-
-st.set_page_config(page_title="EPANET + LLM", layout="wide")
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+pipe = load_llm()
+
 st.title("💧 EPANET + LLM (via WNTR + Falcon LLM)")
 
 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 question:
+if uploaded_file:
     with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as tmp:
         tmp.write(uploaded_file.read())
-        tmp_path = tmp.name
+        inp_path = tmp.name
+    wn = wntr.network.WaterNetworkModel(inp_path)
+
+    question = st.text_input("Ask a question about your water network model")
+
+    if question:
+        with st.spinner("Generating response..."):
+            prompt = f'''
+You are a Python assistant helping with water distribution modeling using WNTR.
+The user has loaded an EPANET model into the variable `wn` using:
+    wn = wntr.network.WaterNetworkModel('model.inp')
+
+Now the user is asking this question:
+    {question}
+
+Write a Python function `def answer():` that computes and returns the result as a variable named `result`.
+
+Use:
+- `wn.pipe_name_list` to access pipe data.
+- `len(wn.pipe_name_list)` for total pipes.
+- Assign output to `result`.
 
-    wn = wntr.network.WaterNetworkModel(tmp_path)
-    sim = wntr.sim.EpanetSimulator(wn)
-    results = sim.run_sim()
+Do not include anything outside the function.
 
-    pipe = load_llm()
-    code = generate_python_code(pipe, question)
+Begin:
+'''
+            llm_output = pipe(prompt, max_new_tokens=200)[0]['generated_text']
+            code_start = llm_output.find("def answer():")
+            code = llm_output[code_start:].strip()
 
-    st.subheader("🧠 Generated Code")
-    st.code(code, language="python")
+            st.markdown("### 🧠 Generated Code")
+            st.code(code, language="python")
 
-    st.subheader("📊 Output")
-    output = run_answer_function(code, wn, results)
-    st.write(output)
+            st.markdown("### 📤 Output")
+            local_vars = {}
+            try:
+                exec(code, {"wn": wn}, local_vars)
+                output = local_vars.get("result", "No result returned.")
+            except Exception as e:
+                output = f"Error in generated code: {e}"
+            st.success(output)
 
-    os.remove(tmp_path)