"""Model training, evaluation, and code generation.""" import io import os import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import pandas as pd from sklearn.ensemble import ( GradientBoostingClassifier, GradientBoostingRegressor, RandomForestClassifier, RandomForestRegressor, ) from sklearn.linear_model import LinearRegression, LogisticRegression from sklearn.metrics import ( accuracy_score, confusion_matrix, f1_score, mean_absolute_error, mean_squared_error, precision_score, r2_score, recall_score, ) from sklearn.model_selection import train_test_split from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor from sklearn.preprocessing import LabelEncoder from sklearn.svm import SVC, SVR from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from data_processor import TEMP_DIR REGRESSION_MODELS = { "Linear Regression": LinearRegression, "Decision Tree Regressor": DecisionTreeRegressor, "Random Forest Regressor": RandomForestRegressor, "Gradient Boosting Regressor": GradientBoostingRegressor, "Support Vector Regressor (SVR)": SVR, "K-Nearest Neighbors Regressor": KNeighborsRegressor, } CLASSIFICATION_MODELS = { "Logistic Regression": LogisticRegression, "Decision Tree Classifier": DecisionTreeClassifier, "Random Forest Classifier": RandomForestClassifier, "Gradient Boosting Classifier": GradientBoostingClassifier, "Support Vector Classifier (SVC)": SVC, "K-Nearest Neighbors Classifier": KNeighborsClassifier, } MODEL_IMPORTS = { "Linear Regression": "from sklearn.linear_model import LinearRegression", "Decision Tree Regressor": "from sklearn.tree import DecisionTreeRegressor", "Random Forest Regressor": "from sklearn.ensemble import RandomForestRegressor", "Gradient Boosting Regressor": "from sklearn.ensemble import GradientBoostingRegressor", "Support Vector Regressor (SVR)": "from sklearn.svm import SVR", "K-Nearest Neighbors Regressor": "from sklearn.neighbors import KNeighborsRegressor", "Logistic Regression": "from sklearn.linear_model import LogisticRegression", "Decision Tree Classifier": "from sklearn.tree import DecisionTreeClassifier", "Random Forest Classifier": "from sklearn.ensemble import RandomForestClassifier", "Gradient Boosting Classifier": "from sklearn.ensemble import GradientBoostingClassifier", "Support Vector Classifier (SVC)": "from sklearn.svm import SVC", "K-Nearest Neighbors Classifier": "from sklearn.neighbors import KNeighborsClassifier", } def suggest_task(df: pd.DataFrame, target: str) -> str: """Suggest 'Classification' or 'Regression' based on the target column.""" s = df[target] if not pd.api.types.is_numeric_dtype(s): return "Classification" if s.nunique() <= 10: return "Classification" return "Regression" def _make_model(name: str, task: str): cls = (CLASSIFICATION_MODELS if task == "Classification" else REGRESSION_MODELS)[name] kwargs = {} if "Logistic" in name: kwargs["max_iter"] = 1000 if "Random Forest" in name or "Gradient Boosting" in name: kwargs["random_state"] = 42 if "Decision Tree" in name: kwargs["random_state"] = 42 return cls(**kwargs) def train_model( df: pd.DataFrame, target: str, model_name: str, task: str, test_size: float = 0.2, ) -> dict: """Train a model and return metrics, plot path, and predictions sample.""" df = df.dropna(subset=[target]) X = df.drop(columns=[target]).select_dtypes(include=np.number) if X.empty: raise ValueError("No numeric feature columns available. Run preprocessing first.") y = df[target] label_encoder = None if task == "Classification" and not pd.api.types.is_numeric_dtype(y): label_encoder = LabelEncoder() y = pd.Series(label_encoder.fit_transform(y), index=df.index) stratify = y if (task == "Classification" and y.value_counts().min() >= 2) else None X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=test_size, random_state=42, stratify=stratify ) model = _make_model(model_name, task) model.fit(X_train, y_train) y_pred = model.predict(X_test) if task == "Classification": avg = "binary" if pd.Series(y).nunique() == 2 else "weighted" metrics = { "Accuracy": round(accuracy_score(y_test, y_pred), 4), "Precision": round(precision_score(y_test, y_pred, average=avg, zero_division=0), 4), "Recall": round(recall_score(y_test, y_pred, average=avg, zero_division=0), 4), "F1 Score": round(f1_score(y_test, y_pred, average=avg, zero_division=0), 4), } plot_path = _plot_confusion(y_test, y_pred, model_name, label_encoder) else: metrics = { "R² Score": round(r2_score(y_test, y_pred), 4), "MAE": round(mean_absolute_error(y_test, y_pred), 4), "RMSE": round(float(np.sqrt(mean_squared_error(y_test, y_pred))), 4), } plot_path = _plot_regression(y_test, y_pred, model_name) importance_path = _plot_importance(model, X.columns, model_name) return { "model_name": model_name, "task": task, "target": target, "features": X.columns.tolist(), "n_train": len(X_train), "n_test": len(X_test), "metrics": metrics, "plot_path": plot_path, "importance_path": importance_path, } def _plot_confusion(y_test, y_pred, model_name, label_encoder=None): cm = confusion_matrix(y_test, y_pred) labels = label_encoder.classes_ if label_encoder is not None else sorted(set(y_test)) fig, ax = plt.subplots(figsize=(5.5, 4.5)) im = ax.imshow(cm, cmap="Blues") ax.set_xticks(range(len(labels)), labels=[str(l) for l in labels], rotation=45, ha="right") ax.set_yticks(range(len(labels)), labels=[str(l) for l in labels]) for i in range(cm.shape[0]): for j in range(cm.shape[1]): ax.text(j, i, cm[i, j], ha="center", va="center", color="white" if cm[i, j] > cm.max() / 2 else "black") ax.set_xlabel("Predicted") ax.set_ylabel("Actual") ax.set_title(f"Confusion Matrix — {model_name}") fig.colorbar(im) fig.tight_layout() path = os.path.join(TEMP_DIR, "result_plot.png") fig.savefig(path, dpi=120) plt.close(fig) return path def _plot_regression(y_test, y_pred, model_name): fig, ax = plt.subplots(figsize=(5.5, 4.5)) ax.scatter(y_test, y_pred, alpha=0.5, edgecolors="none") lims = [min(y_test.min(), y_pred.min()), max(y_test.max(), y_pred.max())] ax.plot(lims, lims, "r--", label="Perfect prediction") ax.set_xlabel("Actual") ax.set_ylabel("Predicted") ax.set_title(f"Actual vs Predicted — {model_name}") ax.legend() fig.tight_layout() path = os.path.join(TEMP_DIR, "result_plot.png") fig.savefig(path, dpi=120) plt.close(fig) return path def _plot_importance(model, feature_names, model_name): importances = None if hasattr(model, "feature_importances_"): importances = model.feature_importances_ elif hasattr(model, "coef_"): coef = model.coef_ importances = np.abs(coef).mean(axis=0) if coef.ndim > 1 else np.abs(coef) if importances is None: return None order = np.argsort(importances)[-15:] fig, ax = plt.subplots(figsize=(6, 4.5)) ax.barh([feature_names[i] for i in order], importances[order], color="#4C72B0") ax.set_title(f"Feature Importance — {model_name}") fig.tight_layout() path = os.path.join(TEMP_DIR, "importance_plot.png") fig.savefig(path, dpi=120) plt.close(fig) return path def generate_model_code(model_name: str, task: str, target: str, test_size: float) -> str: """Return equivalent standalone sklearn code.""" cls = (CLASSIFICATION_MODELS if task == "Classification" else REGRESSION_MODELS)[model_name] kwargs = [] if "Logistic" in model_name: kwargs.append("max_iter=1000") if any(k in model_name for k in ("Random Forest", "Gradient Boosting", "Decision Tree")): kwargs.append("random_state=42") ctor = f"{cls.__name__}({', '.join(kwargs)})" lines = [ "import numpy as np", "import pandas as pd", "from sklearn.model_selection import train_test_split", MODEL_IMPORTS[model_name], ] if task == "Classification": lines += [ "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score", "from sklearn.preprocessing import LabelEncoder", ] else: lines.append( "from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error" ) lines += [ "", "df = pd.read_csv('cleaned_data.csv')", f"target = {target!r}", "X = df.drop(columns=[target]).select_dtypes(include=np.number)", "y = df[target]", ] if task == "Classification": lines += [ "if not pd.api.types.is_numeric_dtype(y):", " y = LabelEncoder().fit_transform(y)", ] lines += [ "", "X_train, X_test, y_train, y_test = train_test_split(", f" X, y, test_size={test_size}, random_state=42)", "", f"model = {ctor}", "model.fit(X_train, y_train)", "y_pred = model.predict(X_test)", "", ] if task == "Classification": lines += [ "print('Accuracy :', accuracy_score(y_test, y_pred))", "print('Precision:', precision_score(y_test, y_pred, average='weighted'))", "print('Recall :', recall_score(y_test, y_pred, average='weighted'))", "print('F1 Score :', f1_score(y_test, y_pred, average='weighted'))", ] else: lines += [ "print('R2 :', r2_score(y_test, y_pred))", "print('MAE :', mean_absolute_error(y_test, y_pred))", "print('RMSE:', np.sqrt(mean_squared_error(y_test, y_pred)))", ] return "\n".join(lines)