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

workbench/core/artifacts/endpoint_core.py

@@ -36,8 +36,8 @@ 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.pytorch_utils import cross_fold_inference as pytorch_cross_fold
-from workbench.utils.chemprop_utils import cross_fold_inference as chemprop_cross_fold
+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
 
 
@@ -389,7 +389,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 +397,26 @@ 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
+        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 and (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...")
@@ -426,14 +433,47 @@ class EndpointCore(Artifact):
             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
-            )
+
+            # Normalize targets to a list for iteration
+            target_list = targets if isinstance(targets, list) else [targets]
+
+            # 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
+
+                description = target_capture_name.replace("_", " ").title()
+
+                # Drop rows with NaN target values for metrics/plots
+                target_df = prediction_df.dropna(subset=[target])
+
+                # 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)
+                elif model.model_type == ModelType.CLASSIFIER:
+                    target_metrics = self.classification_metrics(target, target_df)
+                else:
+                    target_metrics = pd.DataFrame()
+
+                self._capture_inference_results(
+                    target_capture_name,
+                    target_df,
+                    target,
+                    model.model_type,
+                    target_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)
+                metrics = uq_metrics(prediction_df, primary_target)
                 self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
 
         # Return the prediction DataFrame
@@ -453,12 +493,13 @@ class EndpointCore(Artifact):
         model = ModelCore(self.model_name)
 
         # 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 = pytorch_cross_fold(model, nfolds=nfolds)
+            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_cross_fold(model, nfolds=nfolds)
+            cross_fold_metrics, out_of_fold_df = chemprop_pull_cv(model)
         else:
             self.log.error(f"Cross-Fold Inference not supported for Model Framework: {model.model_framework}.")
             return pd.DataFrame()
@@ -475,9 +516,7 @@ 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
@@ -486,7 +525,7 @@ class EndpointCore(Artifact):
 
         # Is this a UQ Model? If so, run full inference and merge the results
         additional_columns = []
-        if model_type == ModelType.UQ_REGRESSOR:
+        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
@@ -495,9 +534,11 @@ class EndpointCore(Artifact):
             # Run inference on the endpoint to get UQ outputs
             uq_df = self.inference(training_df)
 
-            # Identify UQ-specific columns (quantiles and prediction_std)
+            # 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 == "confidence"
+                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
@@ -513,20 +554,42 @@ class EndpointCore(Artifact):
                 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)
+                # Also compute UQ metrics (use first target for multi-target models)
+                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]
+
+        # 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()
+
+            # Drop rows with NaN target values for metrics/plots
+            target_df = out_of_fold_df.dropna(subset=[target])
+
+            # Compute per-target metrics
+            if model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
+                target_metrics = self.regression_metrics(target, target_df)
+            elif model_type == ModelType.CLASSIFIER:
+                target_metrics = self.classification_metrics(target, target_df)
+            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,
+            )
 
-        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:
@@ -736,19 +799,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
@@ -779,26 +842,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)
@@ -808,6 +857,57 @@ class EndpointCore(Artifact):
         model = ModelCore(self.model_name)
         model._load_inference_metrics(capture_name)
 
+    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
+        """
+        # Start with ID column if present
+        output_columns = []
+        if id_column and id_column in pred_results_df.columns:
+            output_columns.append(id_column)
+
+        # Add target column if present
+        if target and target in pred_results_df.columns:
+            output_columns.append(target)
+
+        # Build the output DataFrame
+        output_df = pred_results_df[output_columns].copy() if output_columns else pd.DataFrame()
+
+        # For multi-task: map {target}_pred -> prediction, {target}_pred_std -> prediction_std
+        # For single-task: just grab prediction and prediction_std columns directly
+        pred_col = f"{target}_pred"
+        std_col = f"{target}_pred_std"
+        if pred_col in pred_results_df.columns:
+            # Multi-task columns exist
+            output_df["prediction"] = pred_results_df[pred_col]
+            if std_col in pred_results_df.columns:
+                output_df["prediction_std"] = pred_results_df[std_col]
+        else:
+            # Single-task: grab standard prediction columns
+            for col in ["prediction", "prediction_std"]:
+                if col in pred_results_df.columns:
+                    output_df[col] = pred_results_df[col]
+            # Also grab any _proba columns and UQ columns
+            for col in pred_results_df.columns:
+                if col.endswith("_proba") or col.startswith("q_") or col == "confidence":
+                    output_df[col] = pred_results_df[col]
+
+        # 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(output_df, output_file, index=False)
+
     def regression_metrics(self, target_column: str, prediction_df: pd.DataFrame) -> pd.DataFrame:
         """Compute the performance metrics for this Endpoint
         Args:
@@ -822,24 +922,23 @@ class EndpointCore(Artifact):
             self.log.warning("No predictions were made. Returning empty DataFrame.")
             return pd.DataFrame()
 
+        # Check for prediction column
+        if "prediction" not in prediction_df.columns:
+            self.log.warning("No 'prediction' column found in 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
+        if prediction_df[target_column].isnull().any() or prediction_df["prediction"].isnull().any():
             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])
+            num_nan_prediction = prediction_df["prediction"].isnull().sum()
+            self.log.warning(f"NaNs Found: {target_column} {num_nan_target} and prediction: {num_nan_prediction}.")
+            self.log.warning("Dropping NaN rows for metric computation.")
+            prediction_df = prediction_df.dropna(subset=[target_column, "prediction"])
 
         # Compute the metrics
         try:
             y_true = prediction_df[target_column]
-            y_pred = prediction_df[prediction_col]
+            y_pred = prediction_df["prediction"]
 
             mae = mean_absolute_error(y_true, y_pred)
             rmse = np.sqrt(mean_squared_error(y_true, y_pred))
@@ -871,11 +970,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):
@@ -916,9 +1017,13 @@ class EndpointCore(Artifact):
         Returns:
             pd.DataFrame: DataFrame with the performance metrics
         """
+        # 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 compute metrics on missing predictions)
-        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 metrics calculation")
@@ -938,7 +1043,7 @@ class EndpointCore(Artifact):
         # Calculate precision, recall, f1, and support, handling zero division
         scores = precision_recall_fscore_support(
             prediction_df[target_column],
-            prediction_df[prediction_col],
+            prediction_df["prediction"],
             average=None,
             labels=class_labels,
             zero_division=0,
@@ -986,16 +1091,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()

workbench/core/artifacts/model_core.py

@@ -263,21 +263,25 @@ class ModelCore(Artifact):
         else:
             self.log.important(f"No inference data found for {self.model_name}!")
 
-    def get_inference_metrics(self, capture_name: str = "latest") -> Union[pd.DataFrame, None]:
+    def get_inference_metrics(self, capture_name: str = "any") -> Union[pd.DataFrame, None]:
         """Retrieve the inference performance metrics for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name or "training" (default: "latest")
+            capture_name (str, optional): Specific capture_name (default: "any")
         Returns:
             pd.DataFrame: DataFrame of the Model Metrics
 
         Note:
-            If a capture_name isn't specified this will try to return something reasonable
+            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 == "latest":
-            metrics_df = self.get_inference_metrics("auto_inference")
-            return metrics_df if metrics_df is not None else self.get_inference_metrics("model_training")
+        if capture_name == "any":
+            metric_list = self.list_inference_runs()
+            if metric_list:
+                return self.get_inference_metrics(metric_list[0])
+            else:
+                self.log.warning(f"No performance metrics found for {self.model_name}!")
+                return None
 
         # Grab the metrics captured during model training (could return None)
         if capture_name == "model_training":
@@ -869,7 +873,7 @@ class ModelCore(Artifact):
             return self.df_store.get(f"/workbench/models/{self.name}/shap_data")
         else:
             # Loop over the SHAP data and return a dict of DataFrames
-            shap_dfs = self.df_store.list_subfiles(f"/workbench/models/{self.name}/shap_data")
+            shap_dfs = self.df_store.list(f"/workbench/models/{self.name}/shap_data")
             shap_data = {}
             for df_location in shap_dfs:
                 key = df_location.split("/")[-1]

workbench/core/artifacts/parameter_store_core.py

@@ -0,0 +1,98 @@
+"""ParameterStoreCore: Manages Workbench parameters in a Cloud Based Parameter Store."""
+
+import logging
+
+# Workbench Imports
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Bridges Import
+from workbench_bridges.api import ParameterStore as BridgesParameterStore
+
+
+class ParameterStoreCore(BridgesParameterStore):
+    """ParameterStoreCore: Manages Workbench parameters in a Cloud Based Parameter Store.
+
+    Common Usage:
+        ```python
+        params = ParameterStoreCore()
+
+        # List Parameters
+        params.list()
+
+        ['/workbench/abalone_info',
+         '/workbench/my_data',
+         '/workbench/test',
+         '/workbench/pipelines/my_pipeline']
+
+        # Add Key
+        params.upsert("key", "value")
+        value = params.get("key")
+
+        # Add any data (lists, dictionaries, etc..)
+        my_data = {"key": "value", "number": 4.2, "list": [1,2,3]}
+        params.upsert("my_data", my_data)
+
+        # Retrieve data
+        return_value = params.get("my_data")
+        pprint(return_value)
+
+        {'key': 'value', 'list': [1, 2, 3], 'number': 4.2}
+
+        # Delete parameters
+        param_store.delete("my_data")
+        ```
+    """
+
+    def __init__(self):
+        """ParameterStoreCore Init Method"""
+        session = AWSAccountClamp().boto3_session
+
+        # Initialize parent with workbench config
+        super().__init__(boto3_session=session)
+        self.log = logging.getLogger("workbench")
+
+
+if __name__ == "__main__":
+    """Exercise the ParameterStoreCore Class"""
+
+    # Create a ParameterStoreCore manager
+    param_store = ParameterStoreCore()
+
+    # List the parameters
+    print("Listing Parameters...")
+    print(param_store.list())
+
+    # Add a new parameter
+    param_store.upsert("/workbench/test", "value")
+
+    # Get the parameter
+    print(f"Getting parameter 'test': {param_store.get('/workbench/test')}")
+
+    # Add a dictionary as a parameter
+    sample_dict = {"key": "str_value", "awesome_value": 4.2}
+    param_store.upsert("/workbench/my_data", sample_dict)
+
+    # Retrieve the parameter as a dictionary
+    retrieved_value = param_store.get("/workbench/my_data")
+    print("Retrieved value:", retrieved_value)
+
+    # List the parameters
+    print("Listing Parameters...")
+    print(param_store.list())
+
+    # List the parameters with a prefix
+    print("Listing Parameters with prefix '/workbench':")
+    print(param_store.list("/workbench"))
+
+    # Delete the parameters
+    param_store.delete("/workbench/test")
+    param_store.delete("/workbench/my_data")
+
+    # Out of scope tests
+    param_store.upsert("test", "value")
+    param_store.delete("test")
+
+    # Recursive delete test
+    param_store.upsert("/workbench/test/test1", "value1")
+    param_store.upsert("/workbench/test/test2", "value2")
+    param_store.delete_recursive("workbench/test/")

workbench/core/transforms/features_to_model/features_to_model.py

@@ -1,6 +1,7 @@
 """FeaturesToModel: Train/Create a Model from a Feature Set"""
 
 from pathlib import Path
+from typing import Union
 from sagemaker.estimator import Estimator
 import awswrangler as wr
 from datetime import datetime, timezone
@@ -83,12 +84,17 @@ class FeaturesToModel(Transform):
         self.inference_arch = inference_arch
 
     def transform_impl(
-        self, target_column: str, description: str = None, feature_list: list = None, train_all_data=False, **kwargs
+        self,
+        target_column: Union[str, list[str]],
+        description: str = None,
+        feature_list: list = None,
+        train_all_data=False,
+        **kwargs,
     ):
         """Generic Features to Model: Note you should create a new class and inherit from
         this one to include specific logic for your Feature Set/Model
         Args:
-            target_column (str): Column name of the target variable
+            target_column (str or list[str]): Column name(s) of the target variable(s)
             description (str): Description of the model (optional)
             feature_list (list[str]): A list of columns for the features (default None, will try to guess)
             train_all_data (bool): Train on ALL (100%) of the data (default False)
@@ -105,9 +111,11 @@ class FeaturesToModel(Transform):
         s3_training_path = feature_set.create_s3_training_data()
         self.log.info(f"Created new training data {s3_training_path}...")
 
-        # Report the target column
+        # Report the target column(s)
         self.target_column = target_column
-        self.log.info(f"Target column: {self.target_column}")
+        # Normalize target_column to a list for internal use
+        target_list = [target_column] if isinstance(target_column, str) else (target_column or [])
+        self.log.info(f"Target column(s): {self.target_column}")
 
         # Did they specify a feature list?
         if feature_list:
@@ -134,7 +142,7 @@ class FeaturesToModel(Transform):
                 "is_deleted",
                 "event_time",
                 "training",
-            ] + [self.target_column]
+            ] + target_list
             feature_list = [c for c in all_columns if c not in filter_list]
 
             # AWS Feature Store has 3 user column types (String, Integral, Fractional)
@@ -157,12 +165,14 @@ class FeaturesToModel(Transform):
         self.log.important(f"Feature List for Modeling: {self.model_feature_list}")
 
         # Set up our parameters for the model script
+        # ChemProp expects target_column as a list; other templates expect a string
+        target_for_template = target_list if self.model_framework == ModelFramework.CHEMPROP else self.target_column
         template_params = {
             "model_imports": self.model_import_str,
             "model_type": self.model_type,
             "model_framework": self.model_framework,
             "model_class": self.model_class,
-            "target_column": self.target_column,
+            "target_column": target_for_template,
             "feature_list": self.model_feature_list,
             "compressed_features": feature_set.get_compressed_features(),
             "model_metrics_s3_path": self.model_training_root,
@@ -188,23 +198,27 @@ class FeaturesToModel(Transform):
             # Generate our model script
             script_path = generate_model_script(template_params)
 
-        # Metric Definitions for Regression
+        # Metric Definitions for Regression (matches model script output format)
         if self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]:
             metric_definitions = [
-                {"Name": "RMSE", "Regex": "RMSE: ([0-9.]+)"},
-                {"Name": "MAE", "Regex": "MAE: ([0-9.]+)"},
-                {"Name": "R2", "Regex": "R2: ([0-9.]+)"},
-                {"Name": "NumRows", "Regex": "NumRows: ([0-9]+)"},
+                {"Name": "rmse", "Regex": r"rmse: ([0-9.]+)"},
+                {"Name": "mae", "Regex": r"mae: ([0-9.]+)"},
+                {"Name": "medae", "Regex": r"medae: ([0-9.]+)"},
+                {"Name": "r2", "Regex": r"r2: ([0-9.-]+)"},
+                {"Name": "spearmanr", "Regex": r"spearmanr: ([0-9.-]+)"},
+                {"Name": "support", "Regex": r"support: ([0-9]+)"},
             ]
 
         # Metric Definitions for Classification
         elif self.model_type == ModelType.CLASSIFIER:
             # We need to get creative with the Classification Metrics
+            # Note: Classification only supports single target
+            class_target = target_list[0] if target_list else self.target_column
 
             # Grab all the target column class values (class labels)
             table = feature_set.data_source.table
-            self.class_labels = feature_set.query(f'select DISTINCT {self.target_column} FROM "{table}"')[
-                self.target_column
+            self.class_labels = feature_set.query(f'select DISTINCT {class_target} FROM "{table}"')[
+                class_target
             ].to_list()
 
             # Sanity check on the targets