Update app.py

app.py

@@ -735,6 +735,8 @@ elif app_mode == "Model Training":
                         st.write(f"R-squared: {r2:.4f}")
                     else:
                         from sklearn.metrics import confusion_matrix, roc_curve, auc, precision_recall_curve, classification_report #Import here to avoid library bloat
+                        import seaborn as sns
+                        import matplotlib.pyplot as plt #Added import statement
 
                         #Weighted averaging for metrics for multiclass
                         average_method = "weighted" #changed from None
@@ -763,6 +765,22 @@ elif app_mode == "Model Training":
                         ax_conf.set_title('Confusion Matrix')
                         st.pyplot(fig_conf)
 
+
+                     # Feature Importance (Tree-based Models)
+                    if model_name in ["Random Forest", "Gradient Boosting"] and problem_type == "Classification":
+                        importances = model.feature_importances_ # Assumed tree-based model
+                        feat_importances = pd.Series(importances, index=X_train.columns)
+                        feat_importances = feat_importances.nlargest(20)
+
+                        fig_feat, ax_feat = plt.subplots()
+                        feat_importances.plot(kind='barh', ax=ax_feat)
+                        ax_feat.set_xlabel('Relative Importance')
+                        ax_feat.set_ylabel('Features')
+                        ax_feat.set_title('Feature Importances')
+                        st.pyplot(fig_feat)
+
+                except Exception as e:
+                    st.error(f"An error occurred: {e}")
     else:
         st.write("Please upload and clean data first.")
 
@@ -784,22 +802,21 @@ elif app_mode == "Model Training":
                 st.error(f"Error loading model: {e}")
 
        #Model Evaluation Section
-        if 'X_test' in locals() and st.session_state.model is not None and problem_type == "Regression":
-            y_pred = st.session_state.model.predict(X_test)
-
-            if problem_type == "Regression":
-                mse = mean_squared_error(y_test, y_pred)
-                r2 = r2_score(y_test, y_pred)
-                st.write(f"Mean Squared Error: {mse:.4f}")
-                st.write(f"R-squared: {r2:.4f}")
-            else:
-                from sklearn.metrics import confusion_matrix, roc_curve, auc, precision_recall_curve, classification_report #Import here to avoid library bloat
-
-                #Weighted averaging for metrics for multiclass
-                average_method = "weighted" #changed from None
+        if 'X_test' in locals() and st.session_state.model is not None:
+            try:
+                y_pred = st.session_state.model.predict(X_test)
 
-                accuracy = accuracy_score(y_test, y_pred)
-                st.write(f"Accuracy: {accuracy:.4f}")
+                if problem_type == "Regression":
+                    mse = mean_squared_error(y_test, y_pred)
+                    r2 = r2_score(y_test, y_pred)
+                    st.write(f"Mean Squared Error: {mse:.4f}")
+                    st.write(f"R-squared: {r2:.4f}")
+                else:
+                    from sklearn.metrics import confusion_matrix, roc_curve, auc, precision_recall_curve, classification_report #Import here to avoid library bloat
+                    accuracy = accuracy_score(y_test, y_pred)
+                    st.write(f"Accuracy: {accuracy:.4f}")
+            except Exception as e:
+                 st.error(f"An error occurred during model evaluation: {e}")
 
 elif app_mode == "Predictions":
     st.title("🔮 Make Predictions")