Update app.py

app.py

@@ -1,22 +1,27 @@
 import streamlit as st
 import tempfile
-import os
 import wntr
-from transformers import pipeline
+import os
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
+# Load model once
 @st.cache_resource
 def load_llm():
     model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    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)")
 
 # Upload .inp file
 uploaded_file = st.file_uploader("Upload EPANET .inp file", type=["inp"])
 
-# Initialize the model globally
+# Initialize the WNTR model
 wn = None
 if uploaded_file:
     with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as tmp_file:
@@ -24,32 +29,32 @@ if uploaded_file:
         inp_path = tmp_file.name
     wn = wntr.network.WaterNetworkModel(inp_path)
 
-# Ask question
+# Ask user 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:]
+
+    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
+        else:
+            raise ValueError("Function definition not found in LLM output.")
+
         st.markdown("### 🧠 Generated Code")
         st.code(generated_code, language="python")
 
-        # Execute the function in a controlled environment
         local_env = {"wn": wn}
-        try:
-            exec(generated_code, local_env)
-            result = local_env["answer"]()
-            st.markdown("### 📤 Output")
-            st.success(result)
-        except Exception as e:
-            st.markdown("### 📤 Output")
-            st.error(f"Error in generated code: {e}")
-    else:
-        st.error("Failed to extract Python function from LLM response.")
+        exec(generated_code, local_env)
+        result = local_env["answer"]()
+        st.markdown("### 📤 Output")
+        st.success(result)
+    except Exception as e:
+        st.markdown("### 📤 Output")
+        st.error(f"Error running generated code: {e}")