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

workbench/utils/pytorch_utils.py

@@ -0,0 +1,87 @@
+"""PyTorch Tabular utilities for Workbench models."""
+
+import logging
+import os
+import tarfile
+import tempfile
+from typing import Any, Tuple
+
+import awswrangler as wr
+import pandas as pd
+
+from workbench.utils.aws_utils import pull_s3_data
+from workbench.utils.metrics_utils import compute_metrics_from_predictions
+
+log = logging.getLogger("workbench")
+
+
+def download_and_extract_model(s3_uri: str, model_dir: str) -> None:
+    """Download and extract a PyTorch model artifact from S3.
+
+    Args:
+        s3_uri: S3 URI of the model.tar.gz artifact
+        model_dir: Local directory to extract the model to
+    """
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+        tmp_path = tmp.name
+
+    try:
+        wr.s3.download(path=s3_uri, local_file=tmp_path)
+        with tarfile.open(tmp_path, "r:gz") as tar:
+            tar.extractall(model_dir)
+        log.info(f"Extracted model to {model_dir}")
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+
+
+def pull_cv_results(workbench_model: Any) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """Pull cross-validation results from AWS training artifacts.
+
+    This retrieves the validation predictions saved during model training and
+    computes metrics directly from them. For PyTorch models trained with
+    n_folds > 1, these are out-of-fold predictions from k-fold cross-validation.
+
+    Args:
+        workbench_model: Workbench model object
+
+    Returns:
+        Tuple of:
+            - DataFrame with computed metrics
+            - DataFrame with validation predictions
+    """
+    # Get the validation predictions from S3
+    s3_path = f"{workbench_model.model_training_path}/validation_predictions.csv"
+    predictions_df = pull_s3_data(s3_path)
+
+    if predictions_df is None:
+        raise ValueError(f"No validation predictions found at {s3_path}")
+
+    log.info(f"Pulled {len(predictions_df)} validation predictions from {s3_path}")
+
+    # Compute metrics from predictions
+    target = workbench_model.target()
+    class_labels = workbench_model.class_labels()
+
+    if target in predictions_df.columns and "prediction" in predictions_df.columns:
+        metrics_df = compute_metrics_from_predictions(predictions_df, target, class_labels)
+    else:
+        metrics_df = pd.DataFrame()
+
+    return metrics_df, predictions_df
+
+
+if __name__ == "__main__":
+    from workbench.api import Model
+
+    # Test pulling CV results
+    model_name = "aqsol-reg-pytorch"
+    print(f"Loading Workbench model: {model_name}")
+    model = Model(model_name)
+    print(f"Model Framework: {model.model_framework}")
+
+    # Pull CV results from training artifacts
+    metrics_df, predictions_df = pull_cv_results(model)
+    print(f"\nMetrics:\n{metrics_df}")
+    print(f"\nPredictions shape: {predictions_df.shape}")
+    print(f"Predictions columns: {predictions_df.columns.tolist()}")

workbench/utils/shap_utils.py

@@ -9,6 +9,7 @@ from typing import Optional, List, Tuple, Dict, Union
 from workbench.utils.xgboost_model_utils import xgboost_model_from_s3
 from workbench.utils.model_utils import load_category_mappings_from_s3
 from workbench.utils.pandas_utils import convert_categorical_types
+from workbench.model_script_utils.model_script_utils import decompress_features
 
 # Set up the log
 log = logging.getLogger("workbench")
@@ -111,61 +112,6 @@ def shap_values_data(
         return result_df, feature_df
 
 
-def decompress_features(
-    df: pd.DataFrame, features: List[str], compressed_features: List[str]
-) -> Tuple[pd.DataFrame, List[str]]:
-    """Prepare features for the XGBoost model
-
-    Args:
-        df (pd.DataFrame): The features DataFrame
-        features (List[str]): Full list of feature names
-        compressed_features (List[str]): List of feature names to decompress (bitstrings)
-
-    Returns:
-        pd.DataFrame: DataFrame with the decompressed features
-        List[str]: Updated list of feature names after decompression
-
-    Raises:
-        ValueError: If any missing values are found in the specified features
-    """
-
-    # Check for any missing values in the required features
-    missing_counts = df[features].isna().sum()
-    if missing_counts.any():
-        missing_features = missing_counts[missing_counts > 0]
-        print(
-            f"WARNING: Found missing values in features: {missing_features.to_dict()}. "
-            "WARNING: You might want to remove/replace all NaN values before processing."
-        )
-
-    # Decompress the specified compressed features
-    decompressed_features = features
-    for feature in compressed_features:
-        if (feature not in df.columns) or (feature not in features):
-            print(f"Feature '{feature}' not in the features list, skipping decompression.")
-            continue
-
-        # Remove the feature from the list of features to avoid duplication
-        decompressed_features.remove(feature)
-
-        # Handle all compressed features as bitstrings
-        bit_matrix = np.array([list(bitstring) for bitstring in df[feature]], dtype=np.uint8)
-        prefix = feature[:3]
-
-        # Create all new columns at once - avoids fragmentation
-        new_col_names = [f"{prefix}_{i}" for i in range(bit_matrix.shape[1])]
-        new_df = pd.DataFrame(bit_matrix, columns=new_col_names, index=df.index)
-
-        # Add to features list
-        decompressed_features.extend(new_col_names)
-
-        # Drop original column and concatenate new ones
-        df = df.drop(columns=[feature])
-        df = pd.concat([df, new_df], axis=1)
-
-    return df, decompressed_features
-
-
 def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
     """
     Internal function to calculate SHAP values for Workbench Models.
@@ -212,6 +158,14 @@ def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
         log.error("No XGBoost model found in the artifact.")
         return None, None, None, None
 
+    # Get the booster (SHAP requires the booster, not the sklearn wrapper)
+    if hasattr(xgb_model, "get_booster"):
+        # Full sklearn model - extract the booster
+        booster = xgb_model.get_booster()
+    else:
+        # Already a booster
+        booster = xgb_model
+
     # Load category mappings if available
     category_mappings = load_category_mappings_from_s3(model_artifact_uri)
 
@@ -229,8 +183,8 @@ def _calculate_shap_values(workbench_model, sample_df: pd.DataFrame = None):
     # Create a DMatrix with categorical support
     dmatrix = xgb.DMatrix(X, enable_categorical=True)
 
-    # Use XGBoost's built-in SHAP calculation
-    shap_values = xgb_model.predict(dmatrix, pred_contribs=True, strict_shape=True)
+    # Use XGBoost's built-in SHAP calculation (booster method, not sklearn)
+    shap_values = booster.predict(dmatrix, pred_contribs=True, strict_shape=True)
     features_with_bias = features + ["bias"]
 
     # Now we need to subset the columns based on top 10 SHAP values

workbench/utils/workbench_logging.py

@@ -181,9 +181,6 @@ def logging_setup(color_logs=True):
         log.debug("Debugging enabled via WORKBENCH_DEBUG environment variable.")
     else:
         log.setLevel(logging.INFO)
-        # Note: Not using the ThrottlingFilter for now
-        # throttle_filter = ThrottlingFilter(rate_seconds=5)
-        # handler.addFilter(throttle_filter)
 
     # Suppress specific logger
     logging.getLogger("sagemaker.config").setLevel(logging.WARNING)

workbench/utils/workbench_sqs.py

@@ -12,7 +12,7 @@ class WorkbenchSQS:
         self.log = logging.getLogger("workbench")
         self.queue_url = queue_url
 
-        # Grab a Workbench Session (this allows us to assume the Workbench-ExecutionRole)
+        # Grab a Workbench Session
         self.boto3_session = AWSAccountClamp().boto3_session
         print(self.boto3_session)
 

workbench/utils/xgboost_local_crossfold.py

@@ -0,0 +1,267 @@
+"""XGBoost Local Cross-Fold Validation Utilities
+
+This module contains functions for running cross-validation locally on XGBoost models.
+For most use cases, prefer using pull_cv_results() from xgboost_model_utils.py which
+retrieves the CV results that were saved during training on SageMaker.
+
+These local cross-fold functions are useful for:
+- Re-running CV with different fold counts
+- Leave-one-out cross-validation
+- Custom CV experiments
+"""
+
+import logging
+from typing import Any, Tuple
+
+import numpy as np
+import pandas as pd
+import xgboost as xgb
+from sklearn.model_selection import KFold, StratifiedKFold
+from sklearn.preprocessing import LabelEncoder
+
+from workbench.utils.metrics_utils import compute_metrics_from_predictions
+from workbench.utils.pandas_utils import expand_proba_column
+from workbench.utils.xgboost_model_utils import xgboost_model_from_s3
+
+log = logging.getLogger("workbench")
+
+
+def cross_fold_inference(workbench_model: Any, nfolds: int = 5) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """
+    Performs K-fold cross-validation locally with detailed metrics.
+
+    Note: For most use cases, prefer using pull_cv_results() from xgboost_model_utils.py
+    which retrieves the CV results that were saved during training.
+
+    Args:
+        workbench_model: Workbench model object
+        nfolds: Number of folds for cross-validation (default is 5)
+    Returns:
+        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_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(), 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 = 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()
+    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 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 classifier
+    label_encoder = LabelEncoder() if is_classifier else None
+    if label_encoder:
+        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 if is_classifier else KFold)(n_splits=nfolds, shuffle=True, random_state=42)
+
+    # Initialize predictions DataFrame
+    predictions_df = pd.DataFrame({id_col: ids, target_col: 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_for_cv.iloc[train_idx]
+
+        # 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)
+
+        # 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
+
+    # Expand proba columns for classifiers
+    if is_classifier:
+        predictions_df = expand_proba_column(predictions_df, label_encoder.classes_)
+
+    # Compute metrics from the complete out-of-fold predictions
+    class_labels = list(label_encoder.classes_) if is_classifier else None
+    metrics_df = compute_metrics_from_predictions(predictions_df, target_col, class_labels)
+
+    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 Local Cross-Fold Utilities"""
+    from workbench.api import Model
+    from pprint import pprint
+
+    print("\n=== LOCAL CROSS FOLD REGRESSION EXAMPLE ===")
+    model = Model("abalone-regression")
+    results, df = cross_fold_inference(model)
+    pprint(results)
+    print(df.head())
+
+    print("\n=== LOCAL CROSS FOLD CLASSIFICATION EXAMPLE ===")
+    model = Model("wine-classification")
+    results, df = cross_fold_inference(model)
+    pprint(results)
+    print(df.head())