workbench 0.8.202__py3-none-any.whl → 0.8.220__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
@@ -338,12 +330,8 @@ class EndpointCore(Artifact):
         self.details()
         return True
 
-    def auto_inference(self, capture: bool = False) -> pd.DataFrame:
-        """Run inference on the endpoint using FeatureSet data
-
-        Args:
-            capture (bool, optional): Capture the inference results and metrics (default=False)
-        """
+    def auto_inference(self) -> pd.DataFrame:
+        """Run inference on the endpoint using the test data from the model training view"""
 
         # Sanity Check that we have a model
         model = ModelCore(self.get_input())
@@ -351,22 +339,40 @@ class EndpointCore(Artifact):
             self.log.error("No model found for this endpoint. Returning empty DataFrame.")
             return pd.DataFrame()
 
-        # Now get the FeatureSet and make sure it exists
-        fs = FeatureSetCore(model.get_input())
-        if not fs.exists():
-            self.log.error("No FeatureSet found for this endpoint. Returning empty DataFrame.")
+        # Grab the evaluation data from the Model's training view
+        all_df = model.training_view().pull_dataframe()
+        eval_df = all_df[~all_df["training"]]
+
+        # Remove AWS created columns
+        aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time"]
+        eval_df = eval_df.drop(columns=aws_cols, errors="ignore")
+
+        # Run inference
+        return self.inference(eval_df, "auto_inference")
+
+    def full_inference(self) -> pd.DataFrame:
+        """Run inference on the endpoint using all the data from the model training view"""
+
+        # Sanity Check that we have a model
+        model = ModelCore(self.get_input())
+        if not model.exists():
+            self.log.error("No model found for this endpoint. Returning empty DataFrame.")
             return pd.DataFrame()
 
-        # Grab the evaluation data from the FeatureSet
-        table = model.training_view().table
-        eval_df = fs.query(f'SELECT * FROM "{table}" where training = FALSE')
-        capture_name = "auto_inference" if capture else None
-        return self.inference(eval_df, capture_name, id_column=fs.id_column)
+        # Grab the full data from the Model's training view
+        eval_df = model.training_view().pull_dataframe()
+
+        # Remove AWS created columns
+        aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time"]
+        eval_df = eval_df.drop(columns=aws_cols, errors="ignore")
+
+        # Run inference
+        return self.inference(eval_df, "full_inference")
 
     def inference(
         self, eval_df: pd.DataFrame, capture_name: str = None, id_column: str = None, drop_error_rows: bool = False
     ) -> pd.DataFrame:
-        """Run inference and compute performance metrics with optional capture
+        """Run inference on the Endpoint using the provided DataFrame
 
         Args:
             eval_df (pd.DataFrame): DataFrame to run predictions on (must have superset of features)
@@ -389,7 +395,7 @@ class EndpointCore(Artifact):
         # Grab the model features and target column
         model = ModelCore(self.model_name)
         features = model.features()
-        target_column = model.target()
+        targets = model.target()  # Note: We have multi-target models (so this could be a list)
 
         # Run predictions on the evaluation data
         prediction_df = self._predict(eval_df, features, drop_error_rows)
@@ -397,19 +403,25 @@ class EndpointCore(Artifact):
             self.log.warning("No predictions were made. Returning empty DataFrame.")
             return prediction_df
 
+        # Normalize targets to handle both string and list formats
+        if isinstance(targets, list):
+            primary_target = targets[0] if targets else None
+        else:
+            primary_target = targets
+
         # Sanity Check that the target column is present
-        if target_column and (target_column not in prediction_df.columns):
-            self.log.important(f"Target Column {target_column} not found in prediction_df!")
+        if primary_target not in prediction_df.columns:
+            self.log.important(f"Target Column {primary_target} not found in prediction_df!")
             self.log.important("In order to compute metrics, the target column must be present!")
             metrics = pd.DataFrame()
 
         # Compute the standard performance metrics for this model
         else:
             if model.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
-                prediction_df = self.residuals(target_column, prediction_df)
-                metrics = self.regression_metrics(target_column, prediction_df)
+                prediction_df = self.residuals(primary_target, prediction_df)
+                metrics = self.regression_metrics(primary_target, prediction_df)
             elif model.model_type == ModelType.CLASSIFIER:
-                metrics = self.classification_metrics(target_column, prediction_df)
+                metrics = self.classification_metrics(primary_target, prediction_df)
             else:
                 # For other model types, we don't compute metrics
                 self.log.info(f"Model Type: {model.model_type} doesn't have metrics...")
@@ -420,30 +432,77 @@ class EndpointCore(Artifact):
         print(metrics.head())
 
         # Capture the inference results and metrics
-        if capture_name is not None:
+        if primary_target and capture_name:
 
             # If we don't have an id_column, we'll pull it from the model's FeatureSet
             if id_column is None:
                 fs = FeatureSetCore(model.get_input())
                 id_column = fs.id_column
-            description = capture_name.replace("_", " ").title()
-            self._capture_inference_results(
-                capture_name, prediction_df, target_column, model.model_type, metrics, description, features, id_column
-            )
-
-            # For UQ Models we also capture the uncertainty metrics
-            if model.model_type in [ModelType.UQ_REGRESSOR]:
-                metrics = uq_metrics(prediction_df, target_column)
+
+            # Normalize targets to a list for iteration
+            target_list = targets if isinstance(targets, list) else [targets]
+            primary_target = target_list[0]
+
+            # 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
+
+            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])
+
+                # For multi-target models, prediction column is {target}_pred, otherwise "prediction"
+                pred_col = f"{target}_pred" if is_multi_target else "prediction"
+
+                # Compute per-target metrics
+                if model.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
+                    target_metrics = self.regression_metrics(target, target_df, prediction_col=pred_col)
+                elif model.model_type == ModelType.CLASSIFIER:
+                    target_metrics = self.classification_metrics(target, target_df, prediction_col=pred_col)
+                else:
+                    target_metrics = pd.DataFrame()
+
+                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,
+                    )
+
+            # Capture uncertainty metrics if prediction_std is available (UQ, ChemProp, etc.)
+            if "prediction_std" in prediction_df.columns:
+                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
@@ -455,8 +514,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)
@@ -476,58 +535,74 @@ class EndpointCore(Artifact):
             return out_of_fold_df
 
         # Capture the results
-        capture_name = "full_cross_fold"
-        description = capture_name.replace("_", " ").title()
-        target_column = model.target()
+        targets = model.target()  # Note: We have multi-target models (so this could be a list)
         model_type = model.model_type
 
         # Get the id_column from the model's FeatureSet
         fs = FeatureSetCore(model.get_input())
         id_column = fs.id_column
 
-        # Is this a UQ Model? If so, run full inference and merge the results
-        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...")
+        # Normalize targets to a list for iteration
+        target_list = targets if isinstance(targets, list) else [targets]
+        primary_target = target_list[0]
 
-            # Get the training view dataframe for inference
-            training_df = model.training_view().pull_dataframe()
+        # Collect UQ columns (q_*, confidence) for additional tracking
+        additional_columns = [col for col in out_of_fold_df.columns if col.startswith("q_") or col == "confidence"]
+        if additional_columns:
+            self.log.info(f"UQ columns from training: {', '.join(additional_columns)}")
 
-            # Run inference on the endpoint to get UQ outputs
-            uq_df = self.inference(training_df)
+        # Capture uncertainty metrics if prediction_std is available (UQ, ChemProp, etc.)
+        if "prediction_std" in out_of_fold_df.columns:
+            metrics = uq_metrics(out_of_fold_df, primary_target)
+            self.param_store.upsert(f"/workbench/models/{model.name}/inference/full_cross_fold", metrics)
 
-            # Identify UQ-specific columns (quantiles and prediction_std)
-            uq_columns = [
-                col for col in uq_df.columns if col.startswith("q_") or col == "prediction_std" or col == "confidence"
-            ]
+        # 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
 
-            # Merge UQ columns with out-of-fold predictions
-            if uq_columns:
-                # Keep id_column and UQ columns, drop 'prediction' to avoid conflict when merging
-                uq_df = uq_df[[id_column] + uq_columns]
+        for target in target_list:
+            # Drop rows with NaN target values for metrics/plots
+            target_df = out_of_fold_df.dropna(subset=[target])
 
-                # Drop duplicates in uq_df based on id_column
-                uq_df = uq_df.drop_duplicates(subset=[id_column])
+            # For multi-target models, prediction column is {target}_pred, otherwise "prediction"
+            pred_col = f"{target}_pred" if is_multi_target else "prediction"
 
-                # 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
-                metrics = uq_metrics(out_of_fold_df, target_column)
-                self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
+            # Compute per-target metrics
+            if model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
+                target_metrics = self.regression_metrics(target, target_df, prediction_col=pred_col)
+            elif model_type == ModelType.CLASSIFIER:
+                target_metrics = self.classification_metrics(target, target_df, prediction_col=pred_col)
+            else:
+                target_metrics = pd.DataFrame()
+
+            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,
+                )
 
-        self._capture_inference_results(
-            capture_name,
-            out_of_fold_df,
-            target_column,
-            model_type,
-            cross_fold_metrics,
-            description,
-            features=additional_columns,
-            id_column=id_column,
-        )
         return out_of_fold_df
 
     def fast_inference(self, eval_df: pd.DataFrame, threads: int = 4) -> pd.DataFrame:
@@ -737,19 +812,19 @@ class EndpointCore(Artifact):
         self,
         capture_name: str,
         pred_results_df: pd.DataFrame,
-        target_column: str,
+        target: str,
         model_type: ModelType,
         metrics: pd.DataFrame,
         description: str,
         features: list,
         id_column: str = None,
     ):
-        """Internal: Capture the inference results and metrics to S3
+        """Internal: Capture the inference results and metrics to S3 for a single target
 
         Args:
             capture_name (str): Name of the inference capture
             pred_results_df (pd.DataFrame): DataFrame with the prediction results
-            target_column (str): Name of the target column
+            target (str): Target column name
             model_type (ModelType): Type of the model (e.g. REGRESSOR, CLASSIFIER)
             metrics (pd.DataFrame): DataFrame with the performance metrics
             description (str): Description of the inference results
@@ -780,26 +855,12 @@ class EndpointCore(Artifact):
         self.log.info(f"Writing metrics to {inference_capture_path}/inference_metrics.csv")
         wr.s3.to_csv(metrics, f"{inference_capture_path}/inference_metrics.csv", index=False)
 
-        # Grab the ID column and target column if they are present
-        output_columns = []
-        if id_column and id_column in pred_results_df.columns:
-            output_columns.append(id_column)
-        if target_column in pred_results_df.columns:
-            output_columns.append(target_column)
-
-        # Grab the prediction column, any _proba columns, and UQ columns
-        output_columns += [col for col in pred_results_df.columns if "prediction" in col]
-        output_columns += [col for col in pred_results_df.columns if col.endswith("_proba")]
-        output_columns += [col for col in pred_results_df.columns if col.startswith("q_") or col == "confidence"]
-
-        # Write the predictions to our S3 Model Inference Folder
-        self.log.info(f"Writing predictions to {inference_capture_path}/inference_predictions.csv")
-        subset_df = pred_results_df[output_columns]
-        wr.s3.to_csv(subset_df, f"{inference_capture_path}/inference_predictions.csv", index=False)
+        # Save the inference predictions for this target
+        self._save_target_inference(inference_capture_path, pred_results_df, target, id_column)
 
         # CLASSIFIER: Write the confusion matrix to our S3 Model Inference Folder
         if model_type == ModelType.CLASSIFIER:
-            conf_mtx = self.generate_confusion_matrix(target_column, pred_results_df)
+            conf_mtx = self.generate_confusion_matrix(target, pred_results_df)
             self.log.info(f"Writing confusion matrix to {inference_capture_path}/inference_cm.csv")
             # Note: Unlike other dataframes here, we want to write the index (labels) to the CSV
             wr.s3.to_csv(conf_mtx, f"{inference_capture_path}/inference_cm.csv", index=True)
@@ -809,60 +870,54 @@ class EndpointCore(Artifact):
         model = ModelCore(self.model_name)
         model._load_inference_metrics(capture_name)
 
-    def regression_metrics(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
+    def _save_target_inference(
+        self,
+        inference_capture_path: str,
+        pred_results_df: pd.DataFrame,
+        target: str,
+        id_column: str = None,
+    ):
+        """Save inference results for a single target.
+
+        Args:
+            inference_capture_path (str): S3 path for inference capture
+            pred_results_df (pd.DataFrame): DataFrame with prediction results
+            target (str): Target column name
+            id_column (str, optional): Name of the ID column
+        """
+        cols = pred_results_df.columns
+
+        # Build output columns: id, target, prediction, prediction_std, UQ columns, proba columns
+        output_columns = []
+        if id_column and id_column in cols:
+            output_columns.append(id_column)
+        if target and target in cols:
+            output_columns.append(target)
+
+        output_columns += [c for c in ["prediction", "prediction_std"] if c in cols]
+
+        # Add UQ columns (q_*, confidence) and proba columns
+        output_columns += [c for c in cols if c.startswith("q_") or c == "confidence" or c.endswith("_proba")]
+
+        # Write the predictions to S3
+        output_file = f"{inference_capture_path}/inference_predictions.csv"
+        self.log.info(f"Writing predictions to {output_file}")
+        wr.s3.to_csv(pred_results_df[output_columns], output_file, index=False)
+
+    def regression_metrics(
+        self, target_column: str, prediction_df: pd.DataFrame, prediction_col: str = "prediction"
+    ) -> pd.DataFrame:
         """Compute the performance metrics for this Endpoint
+
         Args:
             target_column (str): Name of the target column
             prediction_df (pd.DataFrame): DataFrame with the prediction results
+            prediction_col (str): Name of the prediction column (default: "prediction")
+
         Returns:
             pd.DataFrame: DataFrame with the performance metrics
         """
-
-        # Sanity Check the prediction DataFrame
-        if prediction_df.empty:
-            self.log.warning("No predictions were made. Returning empty DataFrame.")
-            return pd.DataFrame()
-
-        # Check for NaN values in target or prediction columns
-        prediction_col = "prediction" if "prediction" in prediction_df.columns else "predictions"
-        if prediction_df[target_column].isnull().any() or prediction_df[prediction_col].isnull().any():
-            # Compute the number of NaN values in each column
-            num_nan_target = prediction_df[target_column].isnull().sum()
-            num_nan_prediction = prediction_df[prediction_col].isnull().sum()
-            self.log.warning(
-                f"NaNs Found: {target_column} {num_nan_target} and {prediction_col}: {num_nan_prediction}."
-            )
-            self.log.warning(
-                "NaN values found in target or prediction columns. Dropping NaN rows for metric computation."
-            )
-            prediction_df = prediction_df.dropna(subset=[target_column, prediction_col])
-
-        # Compute the metrics
-        try:
-            y_true = prediction_df[target_column]
-            y_pred = prediction_df[prediction_col]
-
-            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])
-        except Exception as e:
-            self.log.warning(f"Error computing regression metrics: {str(e)}")
-            return pd.DataFrame()
+        return compute_regression_metrics(prediction_df, target_column, prediction_col)
 
     def residuals(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
         """Add the residuals to the prediction DataFrame
@@ -872,11 +927,13 @@ class EndpointCore(Artifact):
         Returns:
             pd.DataFrame: DataFrame with two new columns called 'residuals' and 'residuals_abs'
         """
+        # Check for prediction column
+        if "prediction" not in prediction_df.columns:
+            self.log.warning("No 'prediction' column found. Cannot compute residuals.")
+            return prediction_df
 
-        # Compute the residuals
         y_true = prediction_df[target_column]
-        prediction_col = "prediction" if "prediction" in prediction_df.columns else "predictions"
-        y_pred = prediction_df[prediction_col]
+        y_pred = prediction_df["prediction"]
 
         # Check for classification scenario
         if not pd.api.types.is_numeric_dtype(y_true) or not pd.api.types.is_numeric_dtype(y_pred):
@@ -890,92 +947,22 @@ class EndpointCore(Artifact):
 
         return prediction_df
 
-    @staticmethod
-    def validate_proba_columns(prediction_df: pd.DataFrame, class_labels: list, guessing: bool = False):
-        """Ensure probability columns are correctly aligned with class labels
-
-        Args:
-            prediction_df (pd.DataFrame): DataFrame with the prediction results
-            class_labels (list): List of class labels
-            guessing (bool, optional): Whether we're guessing the class labels. Defaults to False.
-        """
-        proba_columns = [col.replace("_proba", "") for col in prediction_df.columns if col.endswith("_proba")]
-
-        if sorted(class_labels) != sorted(proba_columns):
-            if guessing:
-                raise ValueError(f"_proba columns {proba_columns} != GUESSED class_labels {class_labels}!")
-            else:
-                raise ValueError(f"_proba columns {proba_columns} != class_labels {class_labels}!")
-
-    def classification_metrics(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
+    def classification_metrics(
+        self, target_column: str, prediction_df: pd.DataFrame, prediction_col: str = "prediction"
+    ) -> pd.DataFrame:
         """Compute the performance metrics for this Endpoint
 
         Args:
             target_column (str): Name of the target column
             prediction_df (pd.DataFrame): DataFrame with the prediction results
+            prediction_col (str): Name of the prediction column (default: "prediction")
 
         Returns:
             pd.DataFrame: DataFrame with the performance metrics
         """
-        # Drop rows with NaN predictions (can't compute metrics on missing predictions)
-        prediction_col = "prediction" if "prediction" in prediction_df.columns else "predictions"
-        nan_mask = prediction_df[prediction_col].isna()
-        if nan_mask.any():
-            n_nan = nan_mask.sum()
-            self.log.warning(f"Dropping {n_nan} rows with NaN predictions for metrics calculation")
-            prediction_df = prediction_df[~nan_mask].copy()
-
-        # Get the class labels from the model
+        # Get class labels from the model (metrics_utils will infer if None)
         class_labels = ModelCore(self.model_name).class_labels()
-        if class_labels is None:
-            self.log.warning(
-                "Class labels not found in the model. Guessing class labels from the prediction DataFrame."
-            )
-            class_labels = prediction_df[target_column].unique().tolist()
-            self.validate_proba_columns(prediction_df, class_labels, guessing=True)
-        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_col],
-            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
+        return compute_classification_metrics(prediction_df, target_column, class_labels, prediction_col)
 
     def generate_confusion_matrix(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
         """Compute the confusion matrix for this Endpoint
@@ -987,16 +974,20 @@ class EndpointCore(Artifact):
         Returns:
             pd.DataFrame: DataFrame with the confusion matrix
         """
+        # Check for prediction column
+        if "prediction" not in prediction_df.columns:
+            self.log.warning("No 'prediction' column found in DataFrame")
+            return pd.DataFrame()
+
         # Drop rows with NaN predictions (can't include in confusion matrix)
-        prediction_col = "prediction" if "prediction" in prediction_df.columns else "predictions"
-        nan_mask = prediction_df[prediction_col].isna()
+        nan_mask = prediction_df["prediction"].isna()
         if nan_mask.any():
             n_nan = nan_mask.sum()
             self.log.warning(f"Dropping {n_nan} rows with NaN predictions for confusion matrix")
             prediction_df = prediction_df[~nan_mask].copy()
 
         y_true = prediction_df[target_column]
-        y_pred = prediction_df[prediction_col]
+        y_pred = prediction_df["prediction"]
 
         # Get model class labels
         model_class_labels = ModelCore(self.model_name).class_labels()