Update app.py

app.py

@@ -78,6 +78,11 @@ if uploaded_file is not None:
         X = df_cleaned[features]
         y = df_cleaned[target]
         
+        # Remove rows where target variable y contains NaN
+        df_cleaned = df_cleaned.dropna(subset=[target])
+        X = df_cleaned[features]
+        y = df_cleaned[target]
+        
         # Determine if the target is continuous or categorical
         is_classification = len(y.unique()) <= 10  # If target has fewer than or equal to 10 unique values, treat as classification
         
@@ -154,7 +159,7 @@ if uploaded_file is not None:
             # Download correlation heatmap
             st.subheader("Correlation Heatmap")
             correlation_matrix = df_cleaned.select_dtypes(include=['number']).corr()
-            fig, ax = plt.subplots(figsize=(10, 8))
+            fig, ax = plt.subplots(figsize=(8, 6))
             sns.heatmap(correlation_matrix, annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
             st.pyplot(fig)
             fig.savefig("/tmp/correlation_heatmap.png")