Model Development Pipeline Questions

import numpy as np
from sklearn.base import clone, is_classifier, is_regressor
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, mean_squared_error, r2_score

def k_fold_eval(estimator, X, y, k=5, stratify=True, random_state=None):
    """
    Returns: {
      'folds': [{metric_name: value, ...}, ...],
      'mean': {metric_name: mean, ...},
      'std': {metric_name: std, ...}
    }
    """
    X = np.asarray(X)
    y = np.asarray(y)

    # choose splitter
    if stratify and is_classifier(estimator):
        try:
            splitter = StratifiedKFold(n_splits=k, shuffle=True, random_state=random_state)
        except Exception:
            splitter = KFold(n_splits=k, shuffle=True, random_state=random_state)
    else:
        splitter = KFold(n_splits=k, shuffle=True, random_state=random_state)

    per_fold = []
    for train_idx, test_idx in splitter.split(X, y):
        model = clone(estimator)
        model.fit(X[train_idx], y[train_idx])

        if is_classifier(model):
            y_pred = model.predict(X[test_idx])
            metrics = {
                'accuracy': accuracy_score(y[test_idx], y_pred),
                'f1_macro': f1_score(y[test_idx], y_pred, average='macro')
            }
            # add ROC AUC when possible (binary or probabilistic)
            if len(np.unique(y)) == 2:
                if hasattr(model, "predict_proba"):
                    probs = model.predict_proba(X[test_idx])[:, 1]
                    metrics['roc_auc'] = roc_auc_score(y[test_idx], probs)
                elif hasattr(model, "decision_function"):
                    dec = model.decision_function(X[test_idx])
                    metrics['roc_auc'] = roc_auc_score(y[test_idx], dec)
        elif is_regressor(model):
            y_pred = model.predict(X[test_idx])
            metrics = {
                'mse': mean_squared_error(y[test_idx], y_pred),
                'r2': r2_score(y[test_idx], y_pred)
            }
        else:
            # fallback: treat as regressor-like using predict and MSE
            y_pred = model.predict(X[test_idx])
            metrics = {'mse': mean_squared_error(y[test_idx], y_pred)}

        per_fold.append(metrics)

    # aggregate
    keys = set().union(*[m.keys() for m in per_fold])
    means = {k: np.mean([f.get(k, np.nan) for f in per_fold]) for k in keys}
    stds  = {k: np.std( [f.get(k, np.nan) for f in per_fold], ddof=0) for k in keys}

    return {'folds': per_fold, 'mean': means, 'std': stds}