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

workbench/utils/xgboost_model_utils.py

@@ -1,30 +1,22 @@
 """XGBoost Model Utilities"""
 
+import glob
+import hashlib
 import logging
 import os
-import tempfile
-import tarfile
 import pickle
-import glob
+import tempfile
+from typing import Any, List, Optional, Tuple
+
 import awswrangler as wr
-from typing import Optional, List, Tuple
-import hashlib
+import joblib
 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,
-)
-from sklearn.preprocessing import LabelEncoder
 
 # Workbench Imports
-from workbench.utils.model_utils import load_category_mappings_from_s3
+from workbench.utils.aws_utils import pull_s3_data
+from workbench.utils.metrics_utils import compute_metrics_from_predictions
+from workbench.utils.model_utils import load_category_mappings_from_s3, safe_extract_tarfile
 from workbench.utils.pandas_utils import convert_categorical_types
 
 # Set up the log
@@ -34,14 +26,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 +40,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 +132,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 +141,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,149 +249,45 @@ 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]:
-    """
-    Performs K-fold cross-validation with detailed metrics.
+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 XGBoost models trained with
+    n_folds > 1, these are out-of-fold predictions from k-fold cross-validation.
+
     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 computed metrics
+            - DataFrame with validation predictions
     """
-    from workbench.api import FeatureSet
+    # Get the validation predictions from S3
+    s3_path = f"{workbench_model.model_training_path}/validation_predictions.csv"
+    predictions_df = pull_s3_data(s3_path)
 
-    # 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:
-        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
-    # Prepare data
-    fs = FeatureSet(workbench_model.get_input())
-    df = fs.pull_dataframe()
-    feature_cols = workbench_model.features()
-    # Convert string features to categorical
-    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()]
-
-    # Encode target if it's a classification problem
-    label_encoder = LabelEncoder() if is_classifier else None
-    if label_encoder:
-        y = pd.Series(label_encoder.fit_transform(y), name=workbench_model.target())
-    # 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)
-    )
+    if predictions_df is None:
+        raise ValueError(f"No validation predictions found at {s3_path}")
 
-    fold_results = []
-    all_predictions = []
-    all_actuals = []
-    for fold_idx, (train_idx, val_idx) in enumerate(kfold.split(X, y)):
-        X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
-        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]
-
-        # Train the model
-        xgb_model.fit(X_train, y_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}
-
-        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
-            )
-            fold_metrics.update({"precision": float(scores[0]), "recall": float(scores[1]), "fscore": float(scores[2])})
-        else:
-            fold_metrics.update(
-                {
-                    "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_results.append(fold_metrics)
-    # Calculate overall metrics
-    overall_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(
-            {
-                "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_),
-            }
-        )
+    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:
-        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)),
-            }
-        )
-    # 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,
-    }
+        metrics_df = pd.DataFrame()
+
+    return metrics_df, predictions_df
 
 
 if __name__ == "__main__":
     """Exercise the Model Utilities"""
-    from workbench.api import Model, FeatureSet
-    from pprint import pprint
+    from workbench.api import Model
 
     # Test the XGBoost model loading and feature importance
     model = Model("abalone-regression")
@@ -383,34 +299,26 @@ if __name__ == "__main__":
     model_artifact_uri = model.model_data_url()
     xgb_model = xgboost_model_from_s3(model_artifact_uri)
 
-    # Test with UQ Model
-    uq_model = Model("aqsol-uq")
-    _xgb_model = xgboost_model_from_s3(uq_model.model_data_url())
-
-    # 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:")
-    print(leaf_df)
-
-    # Okay, we're going to copy row 3 and insert it into row 7 to make sure the leaf_hash is the same
-    input_df.iloc[7] = input_df.iloc[3]
-    print("DataFrame with Leaf Hash (3 and 7 should match):")
-    leaf_df = add_leaf_hash(model, input_df)
-    print(leaf_df)
-
-    # Test leaf_stats
-    target_col = "class_number_of_rings"
-    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)
+    # Verify enable_categorical is preserved (for debugging/confidence)
+    print(f"Model parameters: {xgb_model.get_params()}")
+    print(f"enable_categorical: {xgb_model.enable_categorical}")
 
-    print("\n=== CROSS FOLD CLASSIFICATION EXAMPLE ===")
+    print("\n=== PULL CV RESULTS EXAMPLE ===")
+    model = Model("abalone-regression")
+    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()}")
+    print(predictions_df.head())
+
+    # Test on a Classifier model
+    print("\n=== CLASSIFIER MODEL TEST ===")
     model = Model("wine-classification")
-    results = cross_fold_inference(model)
-    pprint(results)
+    features = feature_importance(model)
+    print("Feature Importance:")
+    print(features)
+    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()}")
+    print(predictions_df.head())

workbench/web_interface/components/model_plot.py

@@ -36,10 +36,22 @@ class ModelPlot(ComponentInterface):
             if df is None:
                 return self.display_text("No Data")
 
-            # Calculate the distance from the diagonal for each point
+            # Grab the target(s) for this model
             target = model.target()
+
+            # For multi-task models, match target to inference_run name or default to first
+            if isinstance(target, list):
+                target = next((t for t in target if t in inference_run), target[0])
+
+            # Compute error for coloring
             df["error"] = abs(df["prediction"] - df[target])
-            return ScatterPlot().update_properties(df, color="error", regression_line=True)[0]
+            return ScatterPlot().update_properties(
+                df,
+                color="error",
+                regression_line=True,
+                x=target,
+                y="prediction",
+            )[0]
         else:
             return self.display_text(f"Model Type: {model.model_type}\n\n Awesome Plot Coming Soon!")
 

workbench/web_interface/components/plugin_unit_test.py

@@ -156,10 +156,13 @@ class PluginUnitTest:
         """Run the Dash server for the plugin, handling common errors gracefully."""
         while self.is_port_in_use(self.port):
             log.info(f"Port {self.port} is in use. Trying the next one...")
-            self.port += 1  # Increment the port number until an available one is found
+            self.port += 1
 
         log.info(f"Starting Dash server on port {self.port}...")
-        self.app.run(debug=True, use_reloader=False, port=self.port)
+        try:
+            self.app.run(debug=True, use_reloader=False, port=self.port)
+        except KeyboardInterrupt:
+            log.info("Shutting down Dash server...")
 
     @staticmethod
     def is_port_in_use(port):

workbench/web_interface/components/plugins/dashboard_status.py

@@ -72,7 +72,9 @@ class DashboardStatus(PluginInterface):
             details = "**Redis:** 🔴 Failed to Connect<br>"
 
         # Fill in the license details
-        details += f"**Redis Server:** {config_info['REDIS_HOST']}:{config_info.get('REDIS_PORT', 6379)}<br>"
+        redis_host = config_info.get("REDIS_HOST", "NOT SET")
+        redis_port = config_info.get("REDIS_PORT", "NOT SET")
+        details += f"**Redis Server:** {redis_host}:{redis_port}<br>"
         details += f"**Workbench S3 Bucket:** {config_info['WORKBENCH_BUCKET']}<br>"
         details += f"**Plugin Path:** {config_info.get('WORKBENCH_PLUGINS', 'unknown')}<br>"
         details += f"**Themes Path:** {config_info.get('WORKBENCH_THEMES', 'unknown')}<br>"

workbench/web_interface/components/plugins/generated_compounds.py

@@ -5,7 +5,7 @@ import dash_bootstrap_components as dbc
 
 # Workbench Imports
 from workbench.api.compound import Compound
-from workbench.utils.chem_utils import svg_from_smiles
+from workbench.utils.chem_utils.vis import svg_from_smiles
 from workbench.web_interface.components.plugin_interface import PluginInterface, PluginPage, PluginInputType
 from workbench.utils.theme_manager import ThemeManager
 from workbench.utils.ai_summary import AISummary