Update app.py

app.py

@@ -3,7 +3,6 @@ import pandas as pd
 from pycaret.classification import *
 from pycaret.regression import *
 from pycaret.clustering import *
-from sklearn.model_selection import train_test_split
 import os
 
 # Set page config
@@ -47,23 +46,23 @@ elif app_mode == "Model Training":
     problem_type = st.selectbox("Problem Type", ["Classification", "Regression", "Clustering"])
     target = st.selectbox("Select Target Column", df.columns) if problem_type != "Clustering" else None
 
-    if problem_type == "Clustering":
-        st.info("Clustering does not require a target column. PyCaret will automatically group the data.")
-
     if st.button("Setup PyCaret"):
         with st.spinner("Setting up PyCaret..."):
             if problem_type == "Classification":
-                setup(data=df, session_id=123, verbose=False)
+                setup(data=df, target=target, session_id=123, verbose=False)
                 st.session_state.problem_type = "Classification"
+                st.session_state.setup_complete = True
             elif problem_type == "Regression":
-                setup(data=df, session_id=123, verbose=False)
+                setup(data=df, target=target, session_id=123, verbose=False)
                 st.session_state.problem_type = "Regression"
+                st.session_state.setup_complete = True
             elif problem_type == "Clustering":
                 setup(data=df, session_id=123, verbose=False)
                 st.session_state.problem_type = "Clustering"
+                st.session_state.setup_complete = True
             st.success("PyCaret setup complete! You can now train models.")
 
-    if 'problem_type' in st.session_state:
+    if st.session_state.get("setup_complete", False):
         st.subheader("Train Models")
         if st.button("Compare Models"):
             with st.spinner("Comparing models..."):