Update utils/model_training.py

utils/model_training.py

@@ -12,54 +12,49 @@ def train_all_models(X, y):
     """
     Train all models and return a DataFrame with evaluation metrics.
     """
-    metrics = []
+    model_results = []
+    models = [
+        ("Logistic Regression", LogisticRegression(), 'classification'),
+        ("Random Forest", RandomForestClassifier(), 'classification'),
+        ("SVM", SVC(), 'classification'),
+        ("KNN", KNeighborsClassifier(), 'classification'),
+        ("Decision Tree", DecisionTreeClassifier(), 'classification'),
+        ("Naive Bayes", GaussianNB(), 'classification'),
+        ("Linear Regression", LinearRegression(), 'regression'),
+        ("Random Forest Regressor", RandomForestRegressor(), 'regression'),
+        ("SVR", SVR(), 'regression'),
+        ("KNN Regressor", KNeighborsRegressor(), 'regression'),
+        ("Decision Tree Regressor", DecisionTreeRegressor(), 'regression')
+    ]
     
-    if y.dtype == 'object' or len(y.unique()) <= 10:  # Classification
-        classifiers = {
-            'Logistic Regression': LogisticRegression(max_iter=5000),
-            'Decision Tree': DecisionTreeClassifier(),
-            'Random Forest': RandomForestClassifier(),
-            'Support Vector Machine (SVM)': SVC(),
-            'K-Nearest Neighbors (k-NN)': KNeighborsClassifier(),
-            'Naive Bayes': GaussianNB()
-        }
-
-        for name, classifier in classifiers.items():
-            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)
-            classifier.fit(X_train, y_train)
-            y_pred = classifier.predict(X_test)
-            
-            metrics.append({
-                'Model': name,
-                'Accuracy': round(accuracy_score(y_test, y_pred), 2),
-                'Precision': round(precision_score(y_test, y_pred, zero_division=1, average='macro'), 2),
-                'Recall': round(recall_score(y_test, y_pred, zero_division=1, average='macro'), 2),
-                'F1-Score': round(f1_score(y_test, y_pred, zero_division=1, average='macro'), 2),
-                'Cross-Validated Accuracy': round(cross_val_score(classifier, X, y, cv=5, scoring='accuracy').mean(), 2)
-            })
-
-    else:  # Regression
-        regressors = {
-            'Linear Regression': LinearRegression(),
-            'Decision Tree Regressor': DecisionTreeRegressor(),
-            'Random Forest Regressor': RandomForestRegressor(),
-            'Support Vector Regressor (SVR)': SVR(),
-            'K-Nearest Neighbors Regressor (k-NN)': KNeighborsRegressor()
-        }
-
-        for name, regressor in regressors.items():
-            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
-            regressor.fit(X_train, y_train)
-            y_pred = regressor.predict(X_test)
-            
-            metrics.append({
-                'Model': name,
-                'Mean Squared Error (MSE)': round(mean_squared_error(y_test, y_pred), 2),
-                'Mean Absolute Error (MAE)': round(mean_absolute_error(y_test, y_pred), 2),
-                'R² Score': round(r2_score(y_test, y_pred), 2),
-                'Cross-Validated R² Score': round(cross_val_score(regressor, X, y, cv=5, scoring='r2').mean(), 2)
-            })
-
-    # Return metrics as a DataFrame for easy display
-    metrics_df = pd.DataFrame(metrics)
-    return metrics_df
+    for model_name, model, task_type in models:
+        # Split data into train and test sets
+        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
+        
+        # Train the model
+        model.fit(X_train, y_train)
+        
+        # Make predictions
+        y_pred = model.predict(X_test)
+        
+        # Evaluate based on task type
+        if task_type == 'classification':
+            metrics = {
+                "Model": model_name,
+                "Accuracy": accuracy_score(y_test, y_pred),
+                "Precision": precision_score(y_test, y_pred, average='weighted', zero_division=0),
+                "Recall": recall_score(y_test, y_pred, average='weighted', zero_division=0),
+                "F1 Score": f1_score(y_test, y_pred, average='weighted', zero_division=0)
+            }
+        else:
+            metrics = {
+                "Model": model_name,
+                "MSE": mean_squared_error(y_test, y_pred),
+                "MAE": mean_absolute_error(y_test, y_pred),
+                "R2 Score": r2_score(y_test, y_pred)
+            }
+        model_results.append(metrics)
+    
+    # Convert to DataFrame for display
+    model_results_df = pd.DataFrame(model_results)
+    return model_results_df