workbench 0.8.162__py3-none-any.whl → 0.8.202__py3-none-any.whl

workbench/utils/xgboost_model_utils.py

@@ -3,29 +3,30 @@
 import logging
 import os
 import tempfile
-import tarfile
+import joblib
 import pickle
 import glob
 import awswrangler as wr
-from typing import Optional, List, Tuple
+from typing import Optional, List, Tuple, Any
 import hashlib
 import pandas as pd
 import numpy as np
 import xgboost as xgb
-from typing import Dict, Any
 from sklearn.model_selection import KFold, StratifiedKFold
 from sklearn.metrics import (
     precision_recall_fscore_support,
-    confusion_matrix,
     mean_squared_error,
     mean_absolute_error,
     r2_score,
+    median_absolute_error,
+    roc_auc_score,
 )
+from scipy.stats import spearmanr
 from sklearn.preprocessing import LabelEncoder
 
 # Workbench Imports
-from workbench.utils.model_utils import load_category_mappings_from_s3
-from workbench.utils.pandas_utils import convert_categorical_types
+from workbench.utils.model_utils import load_category_mappings_from_s3, safe_extract_tarfile
+from workbench.utils.pandas_utils import convert_categorical_types, expand_proba_column
 
 # Set up the log
 log = logging.getLogger("workbench")
@@ -34,14 +35,12 @@ log = logging.getLogger("workbench")
 def xgboost_model_from_s3(model_artifact_uri: str):
     """
     Download and extract XGBoost model artifact from S3, then load the model into memory.
-    Handles both direct XGBoost model files and pickled models.
-    Ensures categorical feature support is enabled.
 
     Args:
         model_artifact_uri (str): S3 URI of the model artifact.
 
     Returns:
-        Loaded XGBoost model or None if unavailable.
+        Loaded XGBoost model (XGBClassifier, XGBRegressor, or Booster) or None if unavailable.
     """
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -50,68 +49,90 @@ def xgboost_model_from_s3(model_artifact_uri: str):
         wr.s3.download(path=model_artifact_uri, local_file=local_tar_path)
 
         # Extract tarball
-        with tarfile.open(local_tar_path, "r:gz") as tar:
-            tar.extractall(path=tmpdir, filter="data")
+        safe_extract_tarfile(local_tar_path, tmpdir)
 
         # Define model file patterns to search for (in order of preference)
         patterns = [
-            # Direct XGBoost model files
-            os.path.join(tmpdir, "xgboost-model"),
-            os.path.join(tmpdir, "model"),
-            os.path.join(tmpdir, "*.bin"),
+            # Joblib models (preferred - preserves everything)
+            os.path.join(tmpdir, "*model*.joblib"),
+            os.path.join(tmpdir, "xgb*.joblib"),
+            os.path.join(tmpdir, "**", "*model*.joblib"),
+            os.path.join(tmpdir, "**", "xgb*.joblib"),
+            # Pickle models (also preserves everything)
+            os.path.join(tmpdir, "*model*.pkl"),
+            os.path.join(tmpdir, "xgb*.pkl"),
+            os.path.join(tmpdir, "**", "*model*.pkl"),
+            os.path.join(tmpdir, "**", "xgb*.pkl"),
+            # JSON models (fallback - requires reconstruction)
+            os.path.join(tmpdir, "*model*.json"),
+            os.path.join(tmpdir, "xgb*.json"),
             os.path.join(tmpdir, "**", "*model*.json"),
-            os.path.join(tmpdir, "**", "rmse.json"),
-            # Pickled models
-            os.path.join(tmpdir, "*.pkl"),
-            os.path.join(tmpdir, "**", "*.pkl"),
-            os.path.join(tmpdir, "*.pickle"),
-            os.path.join(tmpdir, "**", "*.pickle"),
+            os.path.join(tmpdir, "**", "xgb*.json"),
         ]
 
         # Try each pattern
         for pattern in patterns:
-            # Use glob to find all matching files
             for model_path in glob.glob(pattern, recursive=True):
-                # Determine file type by extension
+                # Skip files that are clearly not XGBoost models
+                filename = os.path.basename(model_path).lower()
+                if any(skip in filename for skip in ["label_encoder", "scaler", "preprocessor", "transformer"]):
+                    log.debug(f"Skipping non-model file: {model_path}")
+                    continue
+
                 _, ext = os.path.splitext(model_path)
 
                 try:
-                    if ext.lower() in [".pkl", ".pickle"]:
-                        # Handle pickled models
+                    if ext == ".joblib":
+                        model = joblib.load(model_path)
+                        # Verify it's actually an XGBoost model
+                        if isinstance(model, (xgb.XGBClassifier, xgb.XGBRegressor, xgb.Booster)):
+                            log.important(f"Loaded XGBoost model from joblib: {model_path}")
+                            return model
+                        else:
+                            log.debug(f"Skipping non-XGBoost object from {model_path}: {type(model)}")
+
+                    elif ext in [".pkl", ".pickle"]:
                         with open(model_path, "rb") as f:
                             model = pickle.load(f)
-
-                        # Handle different model types
-                        if isinstance(model, xgb.Booster):
-                            log.important(f"Loaded XGBoost Booster from pickle: {model_path}")
+                        # Verify it's actually an XGBoost model
+                        if isinstance(model, (xgb.XGBClassifier, xgb.XGBRegressor, xgb.Booster)):
+                            log.important(f"Loaded XGBoost model from pickle: {model_path}")
                             return model
-                        elif hasattr(model, "get_booster"):
-                            log.important(f"Loaded XGBoost model from pipeline: {model_path}")
-                            booster = model.get_booster()
-                            return booster
-                    else:
-                        # Handle direct XGBoost model files
+                        else:
+                            log.debug(f"Skipping non-XGBoost object from {model_path}: {type(model)}")
+
+                    elif ext == ".json":
+                        # JSON files should be XGBoost models by definition
                         booster = xgb.Booster()
                         booster.load_model(model_path)
-                        log.important(f"Loaded XGBoost model directly: {model_path}")
+                        log.important(f"Loaded XGBoost booster from JSON: {model_path}")
                         return booster
+
                 except Exception as e:
-                    log.info(f"Failed to load model from {model_path}: {e}")
-                    continue  # Try the next file
+                    log.debug(f"Failed to load {model_path}: {e}")
+                    continue
 
-    # If no model found
     log.error("No XGBoost model found in the artifact.")
     return None
 
 
-def feature_importance(workbench_model, importance_type: str = "weight") -> Optional[List[Tuple[str, float]]]:
+def feature_importance(workbench_model, importance_type: str = "gain") -> Optional[List[Tuple[str, float]]]:
     """
     Get sorted feature importances from a Workbench Model object.
 
     Args:
         workbench_model: Workbench model object
-        importance_type: Type of feature importance.
-            Options: 'weight', 'gain', 'cover', 'total_gain', 'total_cover'
+        importance_type: Type of feature importance. Options:
+            - 'gain' (default): Average improvement in loss/objective when feature is used.
+                     Best for understanding predictive power of features.
+            - 'weight': Number of times a feature appears in trees (split count).
+                       Useful for understanding model complexity and feature usage frequency.
+            - 'cover': Average number of samples affected when feature is used.
+                      Shows the relative quantity of observations related to this feature.
+            - 'total_gain': Total improvement in loss/objective across all splits.
+                           Similar to 'gain' but not averaged (can be biased toward frequent features).
+            - 'total_cover': Total number of samples affected across all splits.
+                            Similar to 'cover' but not averaged.
 
     Returns:
         List of tuples (feature, importance) sorted by importance value (descending).
@@ -120,7 +141,8 @@ def feature_importance(workbench_model, importance_type: str = "weight") -> Opti
 
     Note:
         XGBoost's get_score() only returns features with non-zero importance.
-        This function ensures all model features are included in the output.
+        This function ensures all model features are included in the output,
+        adding zero values for features that weren't used in any tree splits.
     """
     model_artifact_uri = workbench_model.model_data_url()
     xgb_model = xgboost_model_from_s3(model_artifact_uri)
@@ -128,11 +150,18 @@ def feature_importance(workbench_model, importance_type: str = "weight") -> Opti
         log.error("No XGBoost model found in the artifact.")
         return None
 
-    # Get feature importances (only non-zero features)
-    importances = xgb_model.get_score(importance_type=importance_type)
+    # Check if we got a full sklearn model or just a booster (for backwards compatibility)
+    if hasattr(xgb_model, "get_booster"):
+        # Full sklearn model - get the booster for feature importance
+        booster = xgb_model.get_booster()
+        all_features = booster.feature_names
+    else:
+        # Already a booster (legacy JSON load)
+        booster = xgb_model
+        all_features = xgb_model.feature_names
 
-    # Get all feature names from the model
-    all_features = xgb_model.feature_names
+    # Get feature importances (only non-zero features)
+    importances = booster.get_score(importance_type=importance_type)
 
     # Create complete importance dict with zeros for missing features
     complete_importances = {feat: importances.get(feat, 0.0) for feat in all_features}
@@ -229,148 +258,326 @@ def leaf_stats(df: pd.DataFrame, target_col: str) -> pd.DataFrame:
     return result_df
 
 
-def cross_fold_inference(workbench_model: Any, nfolds: int = 5) -> Dict[str, Any]:
+def cross_fold_inference(workbench_model: Any, nfolds: int = 5) -> Tuple[pd.DataFrame, pd.DataFrame]:
     """
     Performs K-fold cross-validation with detailed metrics.
     Args:
         workbench_model: Workbench model object
         nfolds: Number of folds for cross-validation (default is 5)
     Returns:
-        Dictionary containing:
-            - folds: Dictionary of formatted strings for each fold
-            - summary_metrics: Summary metrics across folds
-            - overall_metrics: Overall metrics for all folds
+        Tuple of:
+            - DataFrame with per-class metrics (and 'all' row for overall metrics)
+            - DataFrame with columns: id, target, prediction, and *_proba columns (for classifiers)
     """
     from workbench.api import FeatureSet
 
     # Load model
-    model_type = workbench_model.model_type.value
     model_artifact_uri = workbench_model.model_data_url()
-    loaded_booster = xgboost_model_from_s3(model_artifact_uri)
-    if loaded_booster is None:
+    loaded_model = xgboost_model_from_s3(model_artifact_uri)
+    if loaded_model is None:
         log.error("No XGBoost model found in the artifact.")
-        return {}
-    # Create the model wrapper
-    is_classifier = model_type == "classifier"
-    xgb_model = (
-        xgb.XGBClassifier(enable_categorical=True) if is_classifier else xgb.XGBRegressor(enable_categorical=True)
-    )
-    xgb_model._Booster = loaded_booster
+        return pd.DataFrame(), pd.DataFrame()
+
+    # Check if we got a full sklearn model or need to create one
+    if isinstance(loaded_model, (xgb.XGBClassifier, xgb.XGBRegressor)):
+        is_classifier = isinstance(loaded_model, xgb.XGBClassifier)
+
+        # Get the model's hyperparameters and ensure enable_categorical=True
+        params = loaded_model.get_params()
+        params["enable_categorical"] = True
+
+        # Create new model with same params but enable_categorical=True
+        if is_classifier:
+            xgb_model = xgb.XGBClassifier(**params)
+        else:
+            xgb_model = xgb.XGBRegressor(**params)
+
+    elif isinstance(loaded_model, xgb.Booster):
+        # Legacy: got a booster, need to wrap it
+        log.warning("Deprecated: Loaded model is a Booster, wrapping in sklearn model.")
+        is_classifier = workbench_model.model_type.value == "classifier"
+        xgb_model = (
+            xgb.XGBClassifier(enable_categorical=True) if is_classifier else xgb.XGBRegressor(enable_categorical=True)
+        )
+        xgb_model._Booster = loaded_model
+    else:
+        log.error(f"Unexpected model type: {type(loaded_model)}")
+        return pd.DataFrame(), pd.DataFrame()
+
     # Prepare data
     fs = FeatureSet(workbench_model.get_input())
-    df = fs.pull_dataframe()
+    df = workbench_model.training_view().pull_dataframe()
+
+    # Extract sample weights if present
+    sample_weights = df.get("sample_weight")
+    if sample_weights is not None:
+        log.info(f"Using sample weights: min={sample_weights.min():.2f}, max={sample_weights.max():.2f}")
+
+    # Get columns
+    id_col = fs.id_column
+    target_col = workbench_model.target()
     feature_cols = workbench_model.features()
-    # Convert string features to categorical
+    print(f"Target column: {target_col}")
+    print(f"Feature columns: {len(feature_cols)} features")
+
+    # Convert string[python] to object, then to category for XGBoost compatibility
     for col in feature_cols:
-        if df[col].dtype in ["object", "string"]:
-            df[col] = df[col].astype("category")
-    # Split X and y
-    X = df[workbench_model.features()]
-    y = df[workbench_model.target()]
+        if pd.api.types.is_string_dtype(df[col]):
+            df[col] = df[col].astype("object").astype("category")
+
+    X = df[feature_cols]
+    y = df[target_col]
+    ids = df[id_col]
 
-    # Encode target if it's a classification problem
+    # Encode target if classifier
     label_encoder = LabelEncoder() if is_classifier else None
     if label_encoder:
-        y = pd.Series(label_encoder.fit_transform(y), name=workbench_model.target())
+        y_encoded = label_encoder.fit_transform(y)
+        y_for_cv = pd.Series(y_encoded, index=y.index, name=target_col)
+    else:
+        y_for_cv = y
+
     # Prepare KFold
-    kfold = (
-        StratifiedKFold(n_splits=nfolds, shuffle=True, random_state=42)
-        if is_classifier
-        else KFold(n_splits=nfolds, shuffle=True, random_state=42)
-    )
+    kfold = (StratifiedKFold if is_classifier else KFold)(n_splits=nfolds, shuffle=True, random_state=42)
+
+    # Initialize results collection
+    fold_metrics = []
+    predictions_df = pd.DataFrame({id_col: ids, target_col: y})
 
-    fold_results = []
-    all_predictions = []
-    all_actuals = []
-    for fold_idx, (train_idx, val_idx) in enumerate(kfold.split(X, y)):
+    # Perform cross-validation
+    for fold_idx, (train_idx, val_idx) in enumerate(kfold.split(X, y_for_cv), 1):
         X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
-        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
+        y_train, y_val = y_for_cv.iloc[train_idx], y_for_cv.iloc[val_idx]
 
-        # Train the model
-        xgb_model.fit(X_train, y_train)
+        # Get sample weights for training fold
+        weights_train = sample_weights.iloc[train_idx] if sample_weights is not None else None
+
+        # Train and predict
+        xgb_model.fit(X_train, y_train, sample_weight=weights_train)
         preds = xgb_model.predict(X_val)
-        all_predictions.extend(preds)
-        all_actuals.extend(y_val)
 
-        # Calculate metrics for this fold
-        fold_metrics = {"fold": fold_idx + 1}
+        # Store predictions (decode if classifier)
+        val_indices = X_val.index
+        if is_classifier:
+            predictions_df.loc[val_indices, "prediction"] = label_encoder.inverse_transform(preds.astype(int))
+            y_proba = xgb_model.predict_proba(X_val)
+            predictions_df.loc[val_indices, "pred_proba"] = pd.Series(y_proba.tolist(), index=val_indices)
+        else:
+            predictions_df.loc[val_indices, "prediction"] = preds
 
+        # Calculate fold metrics
         if is_classifier:
-            y_val_original = label_encoder.inverse_transform(y_val)
-            preds_original = label_encoder.inverse_transform(preds.astype(int))
-            scores = precision_recall_fscore_support(
-                y_val_original, preds_original, average="weighted", zero_division=0
+            y_val_orig = label_encoder.inverse_transform(y_val)
+            preds_orig = label_encoder.inverse_transform(preds.astype(int))
+
+            # Overall weighted metrics
+            prec, rec, f1, _ = precision_recall_fscore_support(
+                y_val_orig, preds_orig, average="weighted", zero_division=0
+            )
+
+            # Per-class F1
+            prec_per_class, rec_per_class, f1_per_class, _ = precision_recall_fscore_support(
+                y_val_orig, preds_orig, average=None, zero_division=0, labels=label_encoder.classes_
+            )
+
+            # ROC-AUC (overall and per-class)
+            roc_auc_overall = roc_auc_score(y_val, y_proba, multi_class="ovr", average="macro")
+            roc_auc_per_class = roc_auc_score(y_val, y_proba, multi_class="ovr", average=None)
+
+            fold_metrics.append(
+                {
+                    "fold": fold_idx,
+                    "precision": prec,
+                    "recall": rec,
+                    "f1": f1,
+                    "roc_auc": roc_auc_overall,
+                    "precision_per_class": prec_per_class,
+                    "recall_per_class": rec_per_class,
+                    "f1_per_class": f1_per_class,
+                    "roc_auc_per_class": roc_auc_per_class,
+                }
             )
-            fold_metrics.update({"precision": float(scores[0]), "recall": float(scores[1]), "fscore": float(scores[2])})
         else:
-            fold_metrics.update(
+            spearman_corr, _ = spearmanr(y_val, preds)
+            fold_metrics.append(
                 {
-                    "rmse": float(np.sqrt(mean_squared_error(y_val, preds))),
-                    "mae": float(mean_absolute_error(y_val, preds)),
-                    "r2": float(r2_score(y_val, preds)),
+                    "fold": fold_idx,
+                    "rmse": np.sqrt(mean_squared_error(y_val, preds)),
+                    "mae": mean_absolute_error(y_val, preds),
+                    "medae": median_absolute_error(y_val, preds),
+                    "r2": r2_score(y_val, preds),
+                    "spearmanr": spearman_corr,
                 }
             )
 
-        fold_results.append(fold_metrics)
-    # Calculate overall metrics
-    overall_metrics = {}
+    # Calculate summary metrics
+    fold_df = pd.DataFrame(fold_metrics)
+
     if is_classifier:
-        all_actuals_original = label_encoder.inverse_transform(all_actuals)
-        all_predictions_original = label_encoder.inverse_transform(all_predictions)
-        scores = precision_recall_fscore_support(
-            all_actuals_original, all_predictions_original, average="weighted", zero_division=0
-        )
-        overall_metrics.update(
+        # Expand the *_proba columns into separate columns for easier handling
+        predictions_df = expand_proba_column(predictions_df, label_encoder.classes_)
+
+        # Build per-class metrics DataFrame
+        metric_rows = []
+
+        # Per-class rows
+        for idx, class_name in enumerate(label_encoder.classes_):
+            prec_scores = np.array([fold["precision_per_class"][idx] for fold in fold_metrics])
+            rec_scores = np.array([fold["recall_per_class"][idx] for fold in fold_metrics])
+            f1_scores = np.array([fold["f1_per_class"][idx] for fold in fold_metrics])
+            roc_auc_scores = np.array([fold["roc_auc_per_class"][idx] for fold in fold_metrics])
+
+            y_orig = label_encoder.inverse_transform(y_for_cv)
+            support = int((y_orig == class_name).sum())
+
+            metric_rows.append(
+                {
+                    "class": class_name,
+                    "precision": prec_scores.mean(),
+                    "recall": rec_scores.mean(),
+                    "f1": f1_scores.mean(),
+                    "roc_auc": roc_auc_scores.mean(),
+                    "support": support,
+                }
+            )
+
+        # Overall 'all' row
+        metric_rows.append(
             {
-                "precision": float(scores[0]),
-                "recall": float(scores[1]),
-                "fscore": float(scores[2]),
-                "confusion_matrix": confusion_matrix(
-                    all_actuals_original, all_predictions_original, labels=label_encoder.classes_
-                ).tolist(),
-                "label_names": list(label_encoder.classes_),
+                "class": "all",
+                "precision": fold_df["precision"].mean(),
+                "recall": fold_df["recall"].mean(),
+                "f1": fold_df["f1"].mean(),
+                "roc_auc": fold_df["roc_auc"].mean(),
+                "support": len(y_for_cv),
             }
         )
+
+        metrics_df = pd.DataFrame(metric_rows)
+
     else:
-        overall_metrics.update(
-            {
-                "rmse": float(np.sqrt(mean_squared_error(all_actuals, all_predictions))),
-                "mae": float(mean_absolute_error(all_actuals, all_predictions)),
-                "r2": float(r2_score(all_actuals, all_predictions)),
-            }
+        # Regression metrics
+        metrics_df = pd.DataFrame(
+            [
+                {
+                    "rmse": fold_df["rmse"].mean(),
+                    "mae": fold_df["mae"].mean(),
+                    "medae": fold_df["medae"].mean(),
+                    "r2": fold_df["r2"].mean(),
+                    "spearmanr": fold_df["spearmanr"].mean(),
+                    "support": len(y_for_cv),
+                }
+            ]
         )
-    # Calculate summary metrics across folds
-    summary_metrics = {}
-    metrics_to_aggregate = ["precision", "recall", "fscore"] if is_classifier else ["rmse", "mae", "r2"]
-
-    for metric in metrics_to_aggregate:
-        values = [fold[metric] for fold in fold_results]
-        summary_metrics[metric] = f"{float(np.mean(values)):.3f} ±{float(np.std(values)):.3f}"
-    # Format fold results as strings (TBD section)
-    formatted_folds = {}
-    for fold_data in fold_results:
-        fold_key = f"Fold {fold_data['fold']}"
-        if is_classifier:
-            formatted_folds[fold_key] = (
-                f"precision: {fold_data['precision']:.3f}  "
-                f"recall: {fold_data['recall']:.3f}  "
-                f"fscore: {fold_data['fscore']:.3f}"
-            )
-        else:
-            formatted_folds[fold_key] = (
-                f"rmse: {fold_data['rmse']:.3f}  mae: {fold_data['mae']:.3f}  r2: {fold_data['r2']:.3f}"
-            )
-    # Return the results
-    return {
-        "summary_metrics": summary_metrics,
-        # "overall_metrics": overall_metrics,
-        "folds": formatted_folds,
-    }
+
+    return metrics_df, predictions_df
+
+
+def leave_one_out_inference(workbench_model: Any) -> pd.DataFrame:
+    """
+    Performs leave-one-out cross-validation (parallelized).
+    For datasets > 1000 rows, first identifies top 100 worst predictions via 10-fold CV,
+    then performs true leave-one-out on those 100 samples.
+    Each model trains on ALL data except one sample.
+    """
+    from workbench.api import FeatureSet
+    from joblib import Parallel, delayed
+    from tqdm import tqdm
+
+    def train_and_predict_one(model_params, is_classifier, X, y, train_idx, val_idx):
+        """Train on train_idx, predict on val_idx."""
+        model = xgb.XGBClassifier(**model_params) if is_classifier else xgb.XGBRegressor(**model_params)
+        model.fit(X[train_idx], y[train_idx])
+        return model.predict(X[val_idx])[0]
+
+    # Load model and get params
+    model_artifact_uri = workbench_model.model_data_url()
+    loaded_model = xgboost_model_from_s3(model_artifact_uri)
+    if loaded_model is None:
+        log.error("No XGBoost model found in the artifact.")
+        return pd.DataFrame()
+
+    if isinstance(loaded_model, (xgb.XGBClassifier, xgb.XGBRegressor)):
+        is_classifier = isinstance(loaded_model, xgb.XGBClassifier)
+        model_params = loaded_model.get_params()
+    elif isinstance(loaded_model, xgb.Booster):
+        log.warning("Deprecated: Loaded model is a Booster, wrapping in sklearn model.")
+        is_classifier = workbench_model.model_type.value == "classifier"
+        model_params = {"enable_categorical": True}
+    else:
+        log.error(f"Unexpected model type: {type(loaded_model)}")
+        return pd.DataFrame()
+
+    # Load and prepare data
+    fs = FeatureSet(workbench_model.get_input())
+    df = workbench_model.training_view().pull_dataframe()
+    id_col = fs.id_column
+    target_col = workbench_model.target()
+    feature_cols = workbench_model.features()
+
+    # Convert string[python] to object, then to category for XGBoost compatibility
+    # This avoids XGBoost's issue with pandas 2.x string[python] dtype in categorical categories
+    for col in feature_cols:
+        if pd.api.types.is_string_dtype(df[col]):
+            # Double conversion: string[python] -> object -> category
+            df[col] = df[col].astype("object").astype("category")
+
+    # Determine which samples to run LOO on
+    if len(df) > 1000:
+        log.important(f"Dataset has {len(df)} rows. Running 10-fold CV to identify top 1000 worst predictions...")
+        _, predictions_df = cross_fold_inference(workbench_model, nfolds=10)
+        predictions_df["residual_abs"] = np.abs(predictions_df[target_col] - predictions_df["prediction"])
+        worst_samples = predictions_df.nlargest(1000, "residual_abs")
+        worst_ids = worst_samples[id_col].values
+        loo_indices = df[df[id_col].isin(worst_ids)].index.values
+        log.important(f"Running leave-one-out CV on 1000 worst samples. Each model trains on {len(df)-1} rows...")
+    else:
+        log.important(f"Running leave-one-out CV on all {len(df)} samples...")
+        loo_indices = df.index.values
+
+    # Prepare full dataset for training
+    X_full = df[feature_cols].values
+    y_full = df[target_col].values
+
+    # Encode target if classifier
+    label_encoder = LabelEncoder() if is_classifier else None
+    if label_encoder:
+        y_full = label_encoder.fit_transform(y_full)
+
+    # Generate LOO splits
+    splits = []
+    for loo_idx in loo_indices:
+        train_idx = np.delete(np.arange(len(X_full)), loo_idx)
+        val_idx = np.array([loo_idx])
+        splits.append((train_idx, val_idx))
+
+    # Parallel execution
+    predictions = Parallel(n_jobs=4)(
+        delayed(train_and_predict_one)(model_params, is_classifier, X_full, y_full, train_idx, val_idx)
+        for train_idx, val_idx in tqdm(splits, desc="LOO CV")
+    )
+
+    # Build results dataframe
+    predictions_array = np.array(predictions)
+    if label_encoder:
+        predictions_array = label_encoder.inverse_transform(predictions_array.astype(int))
+
+    predictions_df = pd.DataFrame(
+        {
+            id_col: df.loc[loo_indices, id_col].values,
+            target_col: df.loc[loo_indices, target_col].values,
+            "prediction": predictions_array,
+        }
+    )
+
+    predictions_df["residual_abs"] = np.abs(predictions_df[target_col] - predictions_df["prediction"])
+
+    return predictions_df
 
 
 if __name__ == "__main__":
     """Exercise the Model Utilities"""
-    from workbench.api import Model, FeatureSet
+    from workbench.api import Model
     from pprint import pprint
 
     # Test the XGBoost model loading and feature importance
@@ -383,11 +590,28 @@ if __name__ == "__main__":
     model_artifact_uri = model.model_data_url()
     xgb_model = xgboost_model_from_s3(model_artifact_uri)
 
+    # Verify enable_categorical is preserved (for debugging/confidence)
+    print(f"Model parameters: {xgb_model.get_params()}")
+    print(f"enable_categorical: {xgb_model.enable_categorical}")
+
     # Test with UQ Model
     uq_model = Model("aqsol-uq")
     _xgb_model = xgboost_model_from_s3(uq_model.model_data_url())
 
+    print("\n=== CROSS FOLD REGRESSION EXAMPLE ===")
+    model = Model("abalone-regression")
+    results, df = cross_fold_inference(model)
+    pprint(results)
+    print(df.head())
+
+    print("\n=== CROSS FOLD CLASSIFICATION EXAMPLE ===")
+    model = Model("wine-classification")
+    results, df = cross_fold_inference(model)
+    pprint(results)
+    print(df.head())
+
     # Test XGBoost add_leaf_hash
+    """
     input_df = FeatureSet(model.get_input()).pull_dataframe()
     leaf_df = add_leaf_hash(model, input_df)
     print("DataFrame with Leaf Hash:")
@@ -404,13 +628,4 @@ if __name__ == "__main__":
     stats_df = leaf_stats(leaf_df, target_col)
     print("DataFrame with Leaf Statistics:")
     print(stats_df)
-
-    print("\n=== CROSS FOLD REGRESSION EXAMPLE ===")
-    model = Model("abalone-regression")
-    results = cross_fold_inference(model)
-    pprint(results)
-
-    print("\n=== CROSS FOLD CLASSIFICATION EXAMPLE ===")
-    model = Model("wine-classification")
-    results = cross_fold_inference(model)
-    pprint(results)
+    """