Update utils/model_training.py

utils/model_training.py

@@ -1,10 +1,19 @@
-from sklearn.model_selection import cross_val_score
-
-def train_models(df, target, features):
-    X = df[features]
-    y = df[target]
+import pandas as pd
+from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.svm import SVC, SVR
+from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
+from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
+from sklearn.naive_bayes import GaussianNB
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, mean_squared_error, mean_absolute_error, r2_score
 
+def train_all_models(X, y):
+    """
+    Train all models and return a DataFrame with evaluation metrics.
+    """
     metrics = []
+    
     if y.dtype == 'object' or len(y.unique()) <= 10:  # Classification
         classifiers = {
             'Logistic Regression': LogisticRegression(max_iter=5000),
@@ -14,11 +23,12 @@ def train_models(df, target, features):
             '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),
@@ -41,6 +51,7 @@ def train_models(df, target, features):
             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),
@@ -49,5 +60,6 @@ def train_models(df, target, features):
                 '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