Update app.py

app.py

@@ -3,62 +3,58 @@ import pandas as pd
 import traceback
 import ast
 import re
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import pipeline
 
 st.set_page_config(page_title="Auto Dashboard Generator", layout="wide")
-st.title("📊 Auto Dashboard Generator (DeepSeek Edition)")
+st.title("📊 Auto Dashboard Generator (Simplified CodeGen)")
 
 @st.cache_resource(show_spinner=True)
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct")
-    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct")
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+def load_generator():
+    return pipeline("text-generation", model="Salesforce/codegen-350M-mono")
 
-generator = load_model()
+generator = load_generator()
 
 uploaded_file = st.file_uploader("📁 Upload CSV", type="csv")
 prompt = st.text_area(
-    "📝 What dashboard should I build?", 
+    "📝 What dashboard should I build?",
     placeholder="Example: Plot churn rate by contract type and tenure"
 )
 
 if st.button("🚀 Generate Dashboard") and uploaded_file and prompt:
     df = pd.read_csv(uploaded_file)
-    st.subheader("📄 Preview Data")
+    st.subheader("📄 Data Preview")
     st.dataframe(df.head())
 
     sample_data = df.sample(min(len(df), 100)).to_csv(index=False)
 
     gen_prompt = (
-        f"You are a Python expert. Use pandas and plotly inside a Streamlit app.\n"
+        f"You are a Python expert. Use pandas and plotly in Streamlit.\n"
         f"Sample data:\n{sample_data}\n\n"
         f"Task: {prompt}\n\n"
         f"Write Python code starting with `def dashboard():` only.\n"
-        f"Important: Use the dataframe named `df` already loaded; do NOT load any CSV files inside the function."
+        f"Important: Use the existing `df` DataFrame; do NOT load CSV files."
     )
 
-    with st.spinner("Generating Python code..."):
+    with st.spinner("Generating code..."):
         try:
-            output = generator(gen_prompt, max_new_tokens=512, do_sample=False, temperature=0.1)[0]['generated_text']
+            output = generator(gen_prompt, max_new_tokens=512, do_sample=False, temperature=0.1)[0]["generated_text"]
 
-            # Clean the output to extract valid dashboard() function only
             if "def dashboard():" in output:
-                code_part = output.split("def dashboard():", 1)[1]
-
-                code_lines = code_part.splitlines()
-                cleaned_lines = []
-                for line in code_lines:
-                    # Stop if line contains markdown or instruction artifacts
-                    if re.search(r'```|` only\.|Note:|Remember to', line):
-                        break
-                    cleaned_lines.append(line)
-                cleaned_code_body = "\n".join(cleaned_lines)
-                code = "def dashboard():\n" + cleaned_code_body
+                code = "def dashboard():" + output.split("def dashboard():", 1)[1]
             else:
-                st.error("🚨 No dashboard() function generated.")
+                st.error("🚨 No function named `dashboard()` found in generated text.")
                 st.code(output)
                 st.stop()
 
+            # Clean residues
+            code_lines = code.splitlines()
+            cleaned_lines = []
+            for line in code_lines:
+                if re.search(r'```|` only\.|Note:|Remember to', line):
+                    break
+                cleaned_lines.append(line)
+            code = "\n".join(cleaned_lines)
+
             st.subheader("🧠 Generated Code")
             st.code(code, language="python")
 
@@ -71,14 +67,14 @@ if st.button("🚀 Generate Dashboard") and uploaded_file and prompt:
                 st.code(code, language="python")
                 st.stop()
 
-            # Execute code safely, passing df, st, and pd
+            # Execute
             exec_globals = {"df": df, "st": st, "pd": pd}
             exec(code, exec_globals)
 
             if "dashboard" in exec_globals:
                 exec_globals["dashboard"]()
             else:
-                st.warning("⚠️ No dashboard() function found.")
+                st.warning("⚠️ `dashboard()` function didn't run.")
         except Exception as e:
-            st.error("🚨 Generation or execution error:")
+            st.error("🚨 Error during generation or execution:")
             st.text(traceback.format_exc())