Update app.py

app.py

@@ -1,46 +1,26 @@
 import streamlit as st
 import pandas as pd
-from utils.data_cleaning import handle_missing_values, remove_outliers_iqr, cap_extreme_values
-from utils.model_training import encode_categorical, train_and_evaluate
-from utils.visualizations import plot_correlation_heatmap, plot_model_performance, save_plot_as_png
-from io import BytesIO
-from utils.data_cleaning import handle_missing_values, remove_outliers_iqr, cap_extreme_values
-from utils.model_training import encode_categorical, train_and_evaluate
-from utils.visualizations import plot_correlation_heatmap, plot_model_performance, save_plot_as_png
-# Streamlit app title
-st.title("Model Training with Outlier Removal, Metrics, and Correlation Heatmap")
+from data_cleaning import handle_missing_values, remove_outliers_iqr, cap_extreme_values
+from model_training import train_classification_model, train_regression_model
+from visualization import plot_correlation_heatmap, save_plot_as_png
 
 # File uploader
+st.title("Model Training with Metrics and Correlation Heatmap")
 uploaded_file = st.file_uploader("Choose a CSV file", type=["csv"])
 
 if uploaded_file is not None:
-    # Read the uploaded CSV file
     df = pd.read_csv(uploaded_file)
 
-    # Display the dataset
+    # Show the dataset
     st.write("Dataset:")
     st.dataframe(df)
 
-    # Convert categorical (str) data to numerical
-    df = encode_categorical(df)
-
-    # Handle missing values
-    st.write("Handling Missing (Null) Values:")
-    fill_method = st.selectbox("Choose how to handle missing values", ["Drop rows", "Fill with mean/median"])
-    df = handle_missing_values(df, method=fill_method)
-
-    # Remove outliers using the IQR method
-    st.write("Removing Outliers Using IQR:")
-    numeric_cols = df.select_dtypes(include=['float64', 'int64']).columns
-    for col in numeric_cols:
-        df = remove_outliers_iqr(df, col)
-
-    # Capping Extreme Values (based on 5% and 95% percentiles)
-    st.write("Handling Extreme Values (Capping):")
+    # Clean data: Missing values, outliers, and extreme values
+    df = handle_missing_values(df)
+    df = remove_outliers_iqr(df)
     df = cap_extreme_values(df)
 
-    # Display dataset after cleaning
-    st.write("Dataset After Cleaning:")
+    st.write("Cleaned Dataset:")
     st.dataframe(df)
 
     # Add clean data download option
@@ -66,40 +46,18 @@ if uploaded_file is not None:
         mime="image/png"
     )
 
-    # Select target variable
+    # Target and features selection
     target = st.selectbox("Select Target Variable", df.columns)
     features = [col for col in df.columns if col != target]
     X = df[features]
     y = df[target]
 
-    # Model Training and Evaluation
-    if len(y.unique()) > 1:  # Ensure the target variable has at least two unique classes/values
-        model_type = 'classification' if y.dtype == 'object' or len(y.unique()) <= 10 else 'regression'
-        st.subheader(f"{model_type.title()} Model Training")
-
-        train_size = st.slider("Select Training Size", min_value=0.1, max_value=0.9, value=0.8)
-        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=1-train_size, random_state=42)
-
-        metrics = train_and_evaluate(X_train, X_test, y_train, y_test, model_type=model_type)
-
-        # Displaying model performance metrics
-        metrics_df = pd.DataFrame(metrics)
-        st.subheader(f"{model_type.title()} Model Performance Metrics")
+    # Train and evaluate models
+    if y.dtype == 'object' or len(y.unique()) <= 10:  # Classification
+        st.subheader("Classification Model Training")
+        metrics_df = train_classification_model(X, y)
         st.dataframe(metrics_df)
-
-        # Model performance graph
-        st.subheader(f"{model_type.title()} Model Performance Bar Graph")
-        performance_graph = plot_model_performance(metrics_df)
-        st.pyplot(performance_graph)
-
-        # Save performance graph as PNG
-        performance_graph_buf = save_plot_as_png(performance_graph)
-        st.download_button(
-            label=f"Download {model_type.title()} Performance Graph as PNG",
-            data=performance_graph_buf,
-            file_name=f"{model_type}_performance_graph.png",
-            mime="image/png"
-        )
-
-    else:
-        st.error("The target variable must contain at least two unique values for classification or regression.")
+    else:  # Regression
+        st.subheader("Regression Model Training")
+        regression_metrics_df = train_regression_model(X, y)
+        st.dataframe(regression_metrics_df)