Update app.py

app.py

@@ -44,7 +44,12 @@ st.set_page_config(
     page_icon="📈",
     initial_sidebar_state="expanded"
 )
-
+# Initial session state setup (at the top of your script)
+if 'raw_data' not in st.session_state:
+    st.session_state.raw_data = None
+if 'cleaned_data' not in st.session_state:
+    st.session_state.cleaned_data = None
+    
 # Security: Set allowed file types
 ALLOWED_EXTENSIONS = {'csv', 'xlsx', 'parquet', 'feather'}
 MAX_FILE_SIZE_MB = 250  # 250MB limit
@@ -165,6 +170,12 @@ if app_mode == "Data Upload":
                 st.error(f"Error loading file: {str(e)}")
                 st.stop()
 
+                # In your Data Upload section, add this when new data is uploaded
+    if uploaded_file is not None:
+        # Reset models when new data is uploaded
+        st.session_state.model = None
+        st.session_state.preprocessor = None
+
         # Data Health Dashboard
         st.subheader("📊 Data Health Dashboard")
         report = enhanced_quality_report(df)
@@ -752,15 +763,19 @@ elif app_mode == "Model Training":
         Choose from a wide range of algorithms and configurations.
     """)
 
-    # Initialize session state variables if they don't exist
+    # Initialize session state variables
     if 'model' not in st.session_state:
         st.session_state.model = None
     if 'preprocessor' not in st.session_state:
         st.session_state.preprocessor = None
-
-    if 'cleaned_data' not in st.session_state or st.session_state.cleaned_data is None:
-        st.warning("Please clean your data in the Smart Cleaning section first.")
-        st.stop()
+    if 'X_train_selected' not in st.session_state:
+        st.session_state.X_train_selected = None
+    if 'X_test_selected' not in st.session_state:
+        st.session_state.X_test_selected = None
+    if 'y_train' not in st.session_state:
+        st.session_state.y_train = None
+    if 'y_test' not in st.session_state:
+        st.session_state.y_test = None
 
     df = st.session_state.cleaned_data.copy()