workbench 0.8.204__py3-none-any.whl → 0.8.212__py3-none-any.whl

workbench/core/artifacts/endpoint_core.py

@@ -12,16 +12,8 @@ from typing import Union, Optional
 import hashlib
 
 # Model Performance Scores
-from sklearn.metrics import (
-    mean_absolute_error,
-    r2_score,
-    median_absolute_error,
-    roc_auc_score,
-    confusion_matrix,
-    precision_recall_fscore_support,
-    mean_squared_error,
-)
-from sklearn.preprocessing import OneHotEncoder
+from sklearn.metrics import confusion_matrix
+from workbench.utils.metrics_utils import compute_regression_metrics, compute_classification_metrics
 
 # SageMaker Imports
 from sagemaker.serializers import CSVSerializer
@@ -35,7 +27,7 @@ from workbench.utils.endpoint_metrics import EndpointMetrics
 from workbench.utils.cache import Cache
 from workbench.utils.s3_utils import compute_s3_object_hash
 from workbench.utils.model_utils import uq_metrics
-from workbench.utils.xgboost_model_utils import cross_fold_inference as xgboost_cross_fold
+from workbench.utils.xgboost_model_utils import pull_cv_results as xgboost_pull_cv
 from workbench.utils.pytorch_utils import pull_cv_results as pytorch_pull_cv
 from workbench.utils.chemprop_utils import pull_cv_results as chemprop_pull_cv
 from workbench_bridges.endpoints.fast_inference import fast_inference
@@ -397,7 +389,6 @@ class EndpointCore(Artifact):
             self.log.warning("No predictions were made. Returning empty DataFrame.")
             return prediction_df
 
-        # FIXME: Multi-target support - currently uses first target for metrics
         # Normalize targets to handle both string and list formats
         if isinstance(targets, list):
             primary_target = targets[0] if targets else None
@@ -436,19 +427,16 @@ class EndpointCore(Artifact):
 
             # Normalize targets to a list for iteration
             target_list = targets if isinstance(targets, list) else [targets]
+            primary_target = target_list[0]
 
-            # For multi-target models, use target-specific capture names (e.g., auto_target1, auto_target2)
-            # For single-target models, use the original capture name for backward compatibility
-            for target in target_list:
-                # Determine capture name: use prefix for multi-target, original name for single-target
-                if len(target_list) > 1:
-                    prefix = "auto" if "auto" in capture_name else capture_name
-                    target_capture_name = f"{prefix}_{target}"
-                else:
-                    target_capture_name = capture_name
+            # For single-target models (99% of cases), just save with capture_name
+            # For multi-target models, save each as {prefix}_{target} plus primary as capture_name
+            is_multi_target = len(target_list) > 1
 
-                description = target_capture_name.replace("_", " ").title()
+            if is_multi_target:
+                prefix = "auto" if capture_name == "auto_inference" else capture_name
 
+            for target in target_list:
                 # Drop rows with NaN target values for metrics/plots
                 target_df = prediction_df.dropna(subset=[target])
 
@@ -460,30 +448,44 @@ class EndpointCore(Artifact):
                 else:
                     target_metrics = pd.DataFrame()
 
-                self._capture_inference_results(
-                    target_capture_name,
-                    target_df,
-                    target,
-                    model.model_type,
-                    target_metrics,
-                    description,
-                    features,
-                    id_column,
-                )
+                if is_multi_target:
+                    # Multi-target: save as {prefix}_{target}
+                    target_capture_name = f"{prefix}_{target}"
+                    description = target_capture_name.replace("_", " ").title()
+                    self._capture_inference_results(
+                        target_capture_name,
+                        target_df,
+                        target,
+                        model.model_type,
+                        target_metrics,
+                        description,
+                        features,
+                        id_column,
+                    )
+
+                # Save primary target (or single target) with original capture_name
+                if target == primary_target:
+                    self._capture_inference_results(
+                        capture_name,
+                        target_df,
+                        target,
+                        model.model_type,
+                        target_metrics,
+                        capture_name.replace("_", " ").title(),
+                        features,
+                        id_column,
+                    )
 
             # For UQ Models we also capture the uncertainty metrics
-            if model.model_type in [ModelType.UQ_REGRESSOR]:
+            if model.model_type == ModelType.UQ_REGRESSOR:
                 metrics = uq_metrics(prediction_df, primary_target)
                 self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
 
         # Return the prediction DataFrame
         return prediction_df
 
-    def cross_fold_inference(self, nfolds: int = 5) -> pd.DataFrame:
-        """Run cross-fold inference (only works for XGBoost models)
-
-        Args:
-            nfolds (int): Number of folds to use for cross-fold (default: 5)
+    def cross_fold_inference(self) -> pd.DataFrame:
+        """Pull cross-fold inference training results for this Endpoint's model
 
         Returns:
             pd.DataFrame: A DataFrame with cross fold predictions
@@ -495,8 +497,8 @@ class EndpointCore(Artifact):
         # Compute CrossFold (Metrics and Prediction Dataframe)
         # For PyTorch and ChemProp, pull pre-computed CV results from training
         if model.model_framework in [ModelFramework.UNKNOWN, ModelFramework.XGBOOST]:
-            cross_fold_metrics, out_of_fold_df = xgboost_cross_fold(model, nfolds=nfolds)
-        elif model.model_framework == ModelFramework.PYTORCH_TABULAR:
+            cross_fold_metrics, out_of_fold_df = xgboost_pull_cv(model)
+        elif model.model_framework == ModelFramework.PYTORCH:
             cross_fold_metrics, out_of_fold_df = pytorch_pull_cv(model)
         elif model.model_framework == ModelFramework.CHEMPROP:
             cross_fold_metrics, out_of_fold_df = chemprop_pull_cv(model)
@@ -523,51 +525,27 @@ class EndpointCore(Artifact):
         fs = FeatureSetCore(model.get_input())
         id_column = fs.id_column
 
-        # Is this a UQ Model? If so, run full inference and merge the results
+        # For UQ models, get UQ columns from training CV results and compute metrics
+        # Note: XGBoost training now saves all UQ columns (q_*, confidence, prediction_std)
         additional_columns = []
-        if model.model_framework == ModelFramework.XGBOOST and model_type == ModelType.UQ_REGRESSOR:
-            self.log.important("UQ Regressor detected, running full inference to get uncertainty estimates...")
-
-            # Get the training view dataframe for inference
-            training_df = model.training_view().pull_dataframe()
-
-            # Run inference on the endpoint to get UQ outputs
-            uq_df = self.inference(training_df)
-
-            # Identify UQ-specific columns (quantiles, prediction_std, *_pred_std)
-            uq_columns = [
-                col
-                for col in uq_df.columns
-                if col.startswith("q_") or col == "prediction_std" or col.endswith("_pred_std") or col == "confidence"
-            ]
-
-            # Merge UQ columns with out-of-fold predictions
+        if model_type == ModelType.UQ_REGRESSOR:
+            uq_columns = [col for col in out_of_fold_df.columns if col.startswith("q_") or col == "confidence"]
             if uq_columns:
-                # Keep id_column and UQ columns, drop 'prediction' to avoid conflict when merging
-                uq_df = uq_df[[id_column] + uq_columns]
-
-                # Drop duplicates in uq_df based on id_column
-                uq_df = uq_df.drop_duplicates(subset=[id_column])
-
-                # Merge UQ columns into out_of_fold_df
-                out_of_fold_df = pd.merge(out_of_fold_df, uq_df, on=id_column, how="left")
                 additional_columns = uq_columns
-                self.log.info(f"Added UQ columns: {', '.join(additional_columns)}")
-
-                # Also compute UQ metrics (use first target for multi-target models)
+                self.log.info(f"UQ columns from training: {', '.join(uq_columns)}")
                 primary_target = targets[0] if isinstance(targets, list) else targets
                 metrics = uq_metrics(out_of_fold_df, primary_target)
                 self.param_store.upsert(f"/workbench/models/{model.name}/inference/full_cross_fold", metrics)
 
         # Normalize targets to a list for iteration
         target_list = targets if isinstance(targets, list) else [targets]
+        primary_target = target_list[0]
 
-        # For multi-target models, use target-specific capture names (e.g., cv_target1, cv_target2)
-        # For single-target models, use "full_cross_fold" for backward compatibility
-        for target in target_list:
-            capture_name = f"cv_{target}"
-            description = capture_name.replace("_", " ").title()
+        # For single-target models (99% of cases), just save as "full_cross_fold"
+        # For multi-target models, save each as cv_{target} plus primary as "full_cross_fold"
+        is_multi_target = len(target_list) > 1
 
+        for target in target_list:
             # Drop rows with NaN target values for metrics/plots
             target_df = out_of_fold_df.dropna(subset=[target])
 
@@ -579,16 +557,33 @@ class EndpointCore(Artifact):
             else:
                 target_metrics = pd.DataFrame()
 
-            self._capture_inference_results(
-                capture_name,
-                target_df,
-                target,
-                model_type,
-                target_metrics,
-                description,
-                features=additional_columns,
-                id_column=id_column,
-            )
+            if is_multi_target:
+                # Multi-target: save as cv_{target}
+                capture_name = f"cv_{target}"
+                description = capture_name.replace("_", " ").title()
+                self._capture_inference_results(
+                    capture_name,
+                    target_df,
+                    target,
+                    model_type,
+                    target_metrics,
+                    description,
+                    features=additional_columns,
+                    id_column=id_column,
+                )
+
+            # Save primary target (or single target) as "full_cross_fold"
+            if target == primary_target:
+                self._capture_inference_results(
+                    "full_cross_fold",
+                    target_df,
+                    target,
+                    model_type,
+                    target_metrics,
+                    "Full Cross Fold",
+                    features=additional_columns,
+                    id_column=id_column,
+                )
 
         return out_of_fold_df
 
@@ -935,29 +930,9 @@ class EndpointCore(Artifact):
             self.log.warning("Dropping NaN rows for metric computation.")
             prediction_df = prediction_df.dropna(subset=[target_column, "prediction"])
 
-        # Compute the metrics
+        # Compute the metrics using shared utilities
         try:
-            y_true = prediction_df[target_column]
-            y_pred = prediction_df["prediction"]
-
-            mae = mean_absolute_error(y_true, y_pred)
-            rmse = np.sqrt(mean_squared_error(y_true, y_pred))
-            r2 = r2_score(y_true, y_pred)
-            # Mean Absolute Percentage Error
-            mape = np.mean(np.where(y_true != 0, np.abs((y_true - y_pred) / y_true), np.abs(y_true - y_pred))) * 100
-            # Median Absolute Error
-            medae = median_absolute_error(y_true, y_pred)
-
-            # Organize and return the metrics
-            metrics = {
-                "MAE": round(mae, 3),
-                "RMSE": round(rmse, 3),
-                "R2": round(r2, 3),
-                "MAPE": round(mape, 3),
-                "MedAE": round(medae, 3),
-                "NumRows": len(prediction_df),
-            }
-            return pd.DataFrame.from_records([metrics])
+            return compute_regression_metrics(prediction_df, target_column)
         except Exception as e:
             self.log.warning(f"Error computing regression metrics: {str(e)}")
             return pd.DataFrame()
@@ -1040,46 +1015,8 @@ class EndpointCore(Artifact):
         else:
             self.validate_proba_columns(prediction_df, class_labels)
 
-        # Calculate precision, recall, f1, and support, handling zero division
-        scores = precision_recall_fscore_support(
-            prediction_df[target_column],
-            prediction_df["prediction"],
-            average=None,
-            labels=class_labels,
-            zero_division=0,
-        )
-
-        # Identify the probability columns and keep them as a Pandas DataFrame
-        proba_columns = [f"{label}_proba" for label in class_labels]
-        y_score = prediction_df[proba_columns]
-
-        # One-hot encode the true labels using all class labels (fit with class_labels)
-        encoder = OneHotEncoder(categories=[class_labels], sparse_output=False)
-        y_true = encoder.fit_transform(prediction_df[[target_column]])
-
-        # Calculate ROC AUC per label and handle exceptions for missing classes
-        roc_auc_per_label = []
-        for i, label in enumerate(class_labels):
-            try:
-                roc_auc = roc_auc_score(y_true[:, i], y_score.iloc[:, i])
-            except ValueError as e:
-                self.log.warning(f"ROC AUC calculation failed for label {label}.")
-                self.log.warning(f"{str(e)}")
-                roc_auc = 0.0
-            roc_auc_per_label.append(roc_auc)
-
-        # Put the scores into a DataFrame
-        score_df = pd.DataFrame(
-            {
-                target_column: class_labels,
-                "precision": scores[0],
-                "recall": scores[1],
-                "f1": scores[2],
-                "roc_auc": roc_auc_per_label,
-                "support": scores[3],
-            }
-        )
-        return score_df
+        # Compute the metrics using shared utilities (returns per-class + 'all' row)
+        return compute_classification_metrics(prediction_df, target_column, class_labels)
 
     def generate_confusion_matrix(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
         """Compute the confusion matrix for this Endpoint

workbench/core/artifacts/model_core.py

@@ -21,7 +21,7 @@ from workbench.utils.aws_utils import newest_path, pull_s3_data
 from workbench.utils.s3_utils import compute_s3_object_hash
 from workbench.utils.shap_utils import shap_values_data, shap_feature_importance
 from workbench.utils.deprecated_utils import deprecated
-from workbench.utils.model_utils import proximity_model
+from workbench.utils.model_utils import published_proximity_model, get_model_hyperparameters
 
 
 class ModelType(Enum):
@@ -44,7 +44,7 @@ class ModelFramework(Enum):
     SKLEARN = "sklearn"
     XGBOOST = "xgboost"
     LIGHTGBM = "lightgbm"
-    PYTORCH_TABULAR = "pytorch_tabular"
+    PYTORCH = "pytorch"
     CHEMPROP = "chemprop"
     TRANSFORMER = "transformer"
     UNKNOWN = "unknown"
@@ -263,11 +263,11 @@ class ModelCore(Artifact):
         else:
             self.log.important(f"No inference data found for {self.model_name}!")
 
-    def get_inference_metrics(self, capture_name: str = "any") -> Union[pd.DataFrame, None]:
+    def get_inference_metrics(self, capture_name: str = "auto") -> Union[pd.DataFrame, None]:
         """Retrieve the inference performance metrics for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name (default: "any")
+            capture_name (str, optional): Specific capture_name (default: "auto")
         Returns:
             pd.DataFrame: DataFrame of the Model Metrics
 
@@ -275,7 +275,7 @@ class ModelCore(Artifact):
             If a capture_name isn't specified this will try to the 'first' available metrics
         """
         # Try to get the auto_capture 'training_holdout' or the training
-        if capture_name == "any":
+        if capture_name == "auto":
             metric_list = self.list_inference_runs()
             if metric_list:
                 return self.get_inference_metrics(metric_list[0])
@@ -303,11 +303,11 @@ class ModelCore(Artifact):
                 self.log.warning(f"Performance metrics {capture_name} not found for {self.model_name}!")
                 return None
 
-    def confusion_matrix(self, capture_name: str = "latest") -> Union[pd.DataFrame, None]:
+    def confusion_matrix(self, capture_name: str = "auto") -> Union[pd.DataFrame, None]:
         """Retrieve the confusion_matrix for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name or "training" (default: "latest")
+            capture_name (str, optional): Specific capture_name or "training" (default: "auto")
         Returns:
             pd.DataFrame: DataFrame of the Confusion Matrix (might be None)
         """
@@ -319,7 +319,7 @@ class ModelCore(Artifact):
             raise ValueError(error_msg)
 
         # Grab the metrics from the Workbench Metadata (try inference first, then training)
-        if capture_name == "latest":
+        if capture_name == "auto":
             cm = self.confusion_matrix("auto_inference")
             return cm if cm is not None else self.confusion_matrix("model_training")
 
@@ -541,6 +541,17 @@ class ModelCore(Artifact):
         else:
             self.log.error(f"Model {self.model_name} is not a classifier!")
 
+    def summary(self) -> dict:
+        """Summary information about this Model
+
+        Returns:
+            dict: Dictionary of summary information about this Model
+        """
+        self.log.info("Computing Model Summary...")
+        summary = super().summary()
+        summary["hyperparameters"] = get_model_hyperparameters(self)
+        return summary
+
     def details(self) -> dict:
         """Additional Details about this Model
 
@@ -565,6 +576,7 @@ class ModelCore(Artifact):
         details["status"] = self.latest_model["ModelPackageStatus"]
         details["approval_status"] = self.latest_model.get("ModelApprovalStatus", "unknown")
         details["image"] = self.container_image().split("/")[-1]  # Shorten the image uri
+        details["hyperparameters"] = get_model_hyperparameters(self)
 
         # Grab the inference and container info
         inference_spec = self.latest_model["InferenceSpecification"]
@@ -575,16 +587,6 @@ class ModelCore(Artifact):
         details["transform_types"] = inference_spec["SupportedTransformInstanceTypes"]
         details["content_types"] = inference_spec["SupportedContentTypes"]
         details["response_types"] = inference_spec["SupportedResponseMIMETypes"]
-        details["model_metrics"] = self.get_inference_metrics()
-        if self.model_type == ModelType.CLASSIFIER:
-            details["confusion_matrix"] = self.confusion_matrix()
-            details["predictions"] = None
-        elif self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
-            details["confusion_matrix"] = None
-            details["predictions"] = self.get_inference_predictions()
-        else:
-            details["confusion_matrix"] = None
-            details["predictions"] = None
 
         # Grab the inference metadata
         details["inference_meta"] = self.get_inference_metadata()
@@ -904,7 +906,7 @@ class ModelCore(Artifact):
         """
         if prox_model_name is None:
             prox_model_name = self.model_name + "-prox"
-        return proximity_model(self, prox_model_name, track_columns=track_columns)
+        return published_proximity_model(self, prox_model_name, track_columns=track_columns)
 
     def delete(self):
         """Delete the Model Packages and the Model Group"""

workbench/core/transforms/features_to_model/features_to_model.py

@@ -228,7 +228,7 @@ class FeaturesToModel(Transform):
                 raise ValueError(msg)
 
             # Dynamically create the metric definitions
-            metrics = ["precision", "recall", "f1"]
+            metrics = ["precision", "recall", "f1", "support"]
             metric_definitions = []
             for t in self.class_labels:
                 for m in metrics:
@@ -254,7 +254,7 @@ class FeaturesToModel(Transform):
         image = ModelImages.get_image_uri(self.sm_session.boto_region_name, self.training_image)
 
         # Use GPU instance for ChemProp/PyTorch, CPU for others
-        if self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH_TABULAR]:
+        if self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH]:
             train_instance_type = "ml.g6.xlarge"  # NVIDIA L4 GPU, ~$0.80/hr
             self.log.important(f"Using GPU instance {train_instance_type} for {self.model_framework.value}")
         else:

workbench/core/transforms/model_to_endpoint/model_to_endpoint.py

@@ -106,7 +106,7 @@ class ModelToEndpoint(Transform):
             from workbench.api import ModelFramework
 
             self.log.info(f"Model Framework: {workbench_model.model_framework}")
-            if workbench_model.model_framework in [ModelFramework.PYTORCH_TABULAR, ModelFramework.CHEMPROP]:
+            if workbench_model.model_framework in [ModelFramework.PYTORCH, ModelFramework.CHEMPROP]:
                 if mem_size < 4096:
                     self.log.important(
                         f"{workbench_model.model_framework} needs at least 4GB of memory (setting to 4GB)"