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

workbench/utils/metrics_utils.py

@@ -0,0 +1,256 @@
+"""Metrics utilities for computing model performance from predictions."""
+
+import logging
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+from scipy.stats import spearmanr
+from sklearn.metrics import (
+    mean_absolute_error,
+    median_absolute_error,
+    precision_recall_fscore_support,
+    r2_score,
+    roc_auc_score,
+    root_mean_squared_error,
+)
+
+log = logging.getLogger("workbench")
+
+
+def validate_proba_columns(predictions_df: pd.DataFrame, class_labels: List[str], guessing: bool = False) -> bool:
+    """Validate that probability columns match class labels.
+
+    Args:
+        predictions_df: DataFrame with prediction results
+        class_labels: List of class labels
+        guessing: Whether class labels were guessed from data
+
+    Returns:
+        True if validation passes
+
+    Raises:
+        ValueError: If probability columns don't match class labels
+    """
+    proba_columns = [col.replace("_proba", "") for col in predictions_df.columns if col.endswith("_proba")]
+
+    if sorted(class_labels) != sorted(proba_columns):
+        label_type = "GUESSED class_labels" if guessing else "class_labels"
+        raise ValueError(f"_proba columns {proba_columns} != {label_type} {class_labels}!")
+    return True
+
+
+def compute_classification_metrics(
+    predictions_df: pd.DataFrame,
+    target_col: str,
+    class_labels: Optional[List[str]] = None,
+    prediction_col: str = "prediction",
+) -> pd.DataFrame:
+    """Compute classification metrics from a predictions DataFrame.
+
+    Args:
+        predictions_df: DataFrame with target and prediction columns
+        target_col: Name of the target column
+        class_labels: List of class labels in order (if None, inferred from target column)
+        prediction_col: Name of the prediction column (default: "prediction")
+
+    Returns:
+        DataFrame with per-class metrics (precision, recall, f1, roc_auc, support)
+        plus a weighted 'all' row. Returns empty DataFrame if validation fails.
+    """
+    # Validate inputs
+    if predictions_df.empty:
+        log.warning("Empty DataFrame provided. Returning empty metrics.")
+        return pd.DataFrame()
+
+    if prediction_col not in predictions_df.columns:
+        log.warning(f"Prediction column '{prediction_col}' not found in DataFrame. Returning empty metrics.")
+        return pd.DataFrame()
+
+    if target_col not in predictions_df.columns:
+        log.warning(f"Target column '{target_col}' not found in DataFrame. Returning empty metrics.")
+        return pd.DataFrame()
+
+    # Handle NaN predictions
+    df = predictions_df.copy()
+    nan_pred = df[prediction_col].isnull().sum()
+    if nan_pred > 0:
+        log.warning(f"Dropping {nan_pred} rows with NaN predictions.")
+        df = df[~df[prediction_col].isnull()]
+
+    if df.empty:
+        log.warning("No valid rows after dropping NaNs. Returning empty metrics.")
+        return pd.DataFrame()
+
+    # Handle class labels
+    guessing = False
+    if class_labels is None:
+        log.warning("Class labels not provided. Inferring from target column.")
+        class_labels = df[target_col].unique().tolist()
+        guessing = True
+
+    # Validate probability columns if present
+    proba_cols = [col for col in df.columns if col.endswith("_proba")]
+    if proba_cols:
+        validate_proba_columns(df, class_labels, guessing=guessing)
+
+    y_true = df[target_col]
+    y_pred = df[prediction_col]
+
+    # Precision, recall, f1, support per class
+    prec, rec, f1, support = precision_recall_fscore_support(y_true, y_pred, labels=class_labels, zero_division=0)
+
+    # ROC AUC per class (requires probability columns and sorted labels)
+    proba_col_names = [f"{label}_proba" for label in class_labels]
+    if all(col in df.columns for col in proba_col_names):
+        # roc_auc_score requires labels to be sorted, so we sort and reorder results back
+        sorted_labels = sorted(class_labels)
+        sorted_proba_cols = [f"{label}_proba" for label in sorted_labels]
+        y_score_sorted = df[sorted_proba_cols].values
+        roc_auc_sorted = roc_auc_score(y_true, y_score_sorted, labels=sorted_labels, multi_class="ovr", average=None)
+        # Map back to original class_labels order
+        label_to_auc = dict(zip(sorted_labels, roc_auc_sorted))
+        roc_auc = np.array([label_to_auc[label] for label in class_labels])
+    else:
+        roc_auc = np.array([None] * len(class_labels))
+
+    # Build per-class metrics
+    metrics_df = pd.DataFrame(
+        {
+            target_col: class_labels,
+            "precision": prec,
+            "recall": rec,
+            "f1": f1,
+            "roc_auc": roc_auc,
+            "support": support.astype(int),
+        }
+    )
+
+    # Add weighted 'all' row
+    total = support.sum()
+    all_row = {
+        target_col: "all",
+        "precision": (prec * support).sum() / total,
+        "recall": (rec * support).sum() / total,
+        "f1": (f1 * support).sum() / total,
+        "roc_auc": (roc_auc * support).sum() / total if roc_auc[0] is not None else None,
+        "support": int(total),
+    }
+    metrics_df = pd.concat([metrics_df, pd.DataFrame([all_row])], ignore_index=True)
+
+    return metrics_df
+
+
+def compute_regression_metrics(
+    predictions_df: pd.DataFrame,
+    target_col: str,
+    prediction_col: str = "prediction",
+) -> pd.DataFrame:
+    """Compute regression metrics from a predictions DataFrame.
+
+    Args:
+        predictions_df: DataFrame with target and prediction columns
+        target_col: Name of the target column
+        prediction_col: Name of the prediction column (default: "prediction")
+
+    Returns:
+        DataFrame with regression metrics (rmse, mae, medae, r2, spearmanr, support)
+        Returns empty DataFrame if validation fails or no valid data.
+    """
+    # Validate inputs
+    if predictions_df.empty:
+        log.warning("Empty DataFrame provided. Returning empty metrics.")
+        return pd.DataFrame()
+
+    if prediction_col not in predictions_df.columns:
+        log.warning(f"Prediction column '{prediction_col}' not found in DataFrame. Returning empty metrics.")
+        return pd.DataFrame()
+
+    if target_col not in predictions_df.columns:
+        log.warning(f"Target column '{target_col}' not found in DataFrame. Returning empty metrics.")
+        return pd.DataFrame()
+
+    # Handle NaN values
+    df = predictions_df[[target_col, prediction_col]].copy()
+    nan_target = df[target_col].isnull().sum()
+    nan_pred = df[prediction_col].isnull().sum()
+    if nan_target > 0 or nan_pred > 0:
+        log.warning(f"NaNs found: {target_col}={nan_target}, {prediction_col}={nan_pred}. Dropping NaN rows.")
+        df = df.dropna()
+
+    if df.empty:
+        log.warning("No valid rows after dropping NaNs. Returning empty metrics.")
+        return pd.DataFrame()
+
+    y_true = df[target_col].values
+    y_pred = df[prediction_col].values
+
+    return pd.DataFrame(
+        [
+            {
+                "rmse": root_mean_squared_error(y_true, y_pred),
+                "mae": mean_absolute_error(y_true, y_pred),
+                "medae": median_absolute_error(y_true, y_pred),
+                "r2": r2_score(y_true, y_pred),
+                "spearmanr": spearmanr(y_true, y_pred).correlation,
+                "support": len(y_true),
+            }
+        ]
+    )
+
+
+def compute_metrics_from_predictions(
+    predictions_df: pd.DataFrame,
+    target_col: str,
+    class_labels: Optional[List[str]] = None,
+    prediction_col: str = "prediction",
+) -> pd.DataFrame:
+    """Compute metrics from a predictions DataFrame.
+
+    Automatically determines if this is classification or regression based on
+    whether class_labels is provided.
+
+    Args:
+        predictions_df: DataFrame with target and prediction columns
+        target_col: Name of the target column
+        class_labels: List of class labels for classification (None for regression)
+        prediction_col: Name of the prediction column (default: "prediction")
+
+    Returns:
+        DataFrame with computed metrics
+    """
+    if target_col not in predictions_df.columns:
+        raise ValueError(f"Target column '{target_col}' not found in predictions DataFrame")
+    if prediction_col not in predictions_df.columns:
+        raise ValueError(f"Prediction column '{prediction_col}' not found in predictions DataFrame")
+
+    if class_labels:
+        return compute_classification_metrics(predictions_df, target_col, class_labels, prediction_col)
+    else:
+        return compute_regression_metrics(predictions_df, target_col, prediction_col)
+
+
+if __name__ == "__main__":
+    # Test with sample data
+    print("Testing classification metrics...")
+    class_df = pd.DataFrame(
+        {
+            "target": ["a", "b", "c", "a", "b", "c", "a", "b", "c", "a"],
+            "prediction": ["a", "b", "c", "a", "b", "a", "a", "b", "c", "b"],
+            "a_proba": [0.8, 0.1, 0.1, 0.7, 0.2, 0.4, 0.9, 0.1, 0.1, 0.3],
+            "b_proba": [0.1, 0.8, 0.1, 0.2, 0.7, 0.3, 0.05, 0.8, 0.2, 0.6],
+            "c_proba": [0.1, 0.1, 0.8, 0.1, 0.1, 0.3, 0.05, 0.1, 0.7, 0.1],
+        }
+    )
+    metrics = compute_metrics_from_predictions(class_df, "target", ["a", "b", "c"])
+    print(metrics.to_string(index=False))
+
+    print("\nTesting regression metrics...")
+    reg_df = pd.DataFrame(
+        {
+            "target": [1.0, 2.0, 3.0, 4.0, 5.0],
+            "prediction": [1.1, 2.2, 2.9, 4.1, 4.8],
+        }
+    )
+    metrics = compute_metrics_from_predictions(reg_df, "target")
+    print(metrics.to_string(index=False))

workbench/utils/model_utils.py

@@ -93,16 +93,17 @@ def get_custom_script_path(package: str, script_name: str) -> Path:
     return script_path
 
 
-def proximity_model_local(model: "Model"):
-    """Create a Proximity Model for this Model
+def proximity_model_local(model: "Model", include_all_columns: bool = False):
+    """Create a FeatureSpaceProximity Model for this Model
 
     Args:
         model (Model): The Model/FeatureSet used to create the proximity model
+        include_all_columns (bool): Include all DataFrame columns in neighbor results (default: False)
 
     Returns:
-        Proximity: The proximity model
+        FeatureSpaceProximity: The proximity model
     """
-    from workbench.algorithms.dataframe.proximity import Proximity  # noqa: F401 (avoid circular import)
+    from workbench.algorithms.dataframe.feature_space_proximity import FeatureSpaceProximity  # noqa: F401
     from workbench.api import Model, FeatureSet  # noqa: F401 (avoid circular import)
 
     # Get Feature and Target Columns from the existing given Model
@@ -121,74 +122,154 @@ def proximity_model_local(model: "Model"):
     model_ids = set(model_df[id_column])
     full_df["in_model"] = full_df[id_column].isin(model_ids)
 
-    # Create and return the Proximity Model
-    return Proximity(full_df, id_column, features, target, track_columns=features)
+    # Create and return the FeatureSpaceProximity Model
+    return FeatureSpaceProximity(
+        full_df, id_column=id_column, features=features, target=target, include_all_columns=include_all_columns
+    )
 
 
-def proximity_model(model: "Model", prox_model_name: str, track_columns: list = None) -> "Model":
-    """Create a proximity model based on the given model
+def fingerprint_prox_model_local(
+    model: "Model",
+    include_all_columns: bool = False,
+    radius: int = 2,
+    n_bits: int = 1024,
+    counts: bool = False,
+):
+    """Create a FingerprintProximity Model for this Model
 
     Args:
-        model (Model): The model to create the proximity model from
-        prox_model_name (str): The name of the proximity model to create
-        track_columns (list, optional): List of columns to track in the proximity model
+        model (Model): The Model used to create the fingerprint proximity model
+        include_all_columns (bool): Include all DataFrame columns in neighbor results (default: False)
+        radius (int): Morgan fingerprint radius (default: 2)
+        n_bits (int): Number of bits for the fingerprint (default: 1024)
+        counts (bool): Use count fingerprints instead of binary (default: False)
+
     Returns:
-        Model: The proximity model
+        FingerprintProximity: The fingerprint proximity model
     """
-    from workbench.api import Model, ModelType, FeatureSet  # noqa: F401 (avoid circular import)
+    from workbench.algorithms.dataframe.fingerprint_proximity import FingerprintProximity  # noqa: F401
+    from workbench.api import Model, FeatureSet  # noqa: F401 (avoid circular import)
 
-    # Get the custom script path for the proximity model
-    script_path = get_custom_script_path("proximity", "feature_space_proximity.template")
+    # Get Target Column from the existing given Model
+    target = model.target()
+
+    # Backtrack our FeatureSet to get the ID column
+    fs = FeatureSet(model.get_input())
+    id_column = fs.id_column
+
+    # Create the Proximity Model from both the full FeatureSet and the Model training data
+    full_df = fs.pull_dataframe()
+    model_df = model.training_view().pull_dataframe()
+
+    # Mark rows that are in the model
+    model_ids = set(model_df[id_column])
+    full_df["in_model"] = full_df[id_column].isin(model_ids)
+
+    # Create and return the FingerprintProximity Model
+    return FingerprintProximity(
+        full_df,
+        id_column=id_column,
+        target=target,
+        include_all_columns=include_all_columns,
+        radius=radius,
+        n_bits=n_bits,
+    )
+
+
+def noise_model_local(model: "Model"):
+    """Create a NoiseModel for detecting noisy/problematic samples in a Model's training data.
+
+    Args:
+        model (Model): The Model used to create the noise model
+
+    Returns:
+        NoiseModel: The noise model with precomputed noise scores for all samples
+    """
+    from workbench.algorithms.models.noise_model import NoiseModel  # noqa: F401 (avoid circular import)
+    from workbench.api import Model, FeatureSet  # noqa: F401 (avoid circular import)
 
     # Get Feature and Target Columns from the existing given Model
     features = model.features()
     target = model.target()
 
-    # Create the Proximity Model from our FeatureSet
+    # Backtrack our FeatureSet to get the ID column
     fs = FeatureSet(model.get_input())
-    prox_model = fs.to_model(
-        name=prox_model_name,
-        model_type=ModelType.PROXIMITY,
-        feature_list=features,
-        target_column=target,
-        description=f"Proximity Model for {model.name}",
-        tags=["proximity", model.name],
-        custom_script=script_path,
-        custom_args={"track_columns": track_columns},
-    )
-    return prox_model
+    id_column = fs.id_column
+
+    # Create the NoiseModel from both the full FeatureSet and the Model training data
+    full_df = fs.pull_dataframe()
+    model_df = model.training_view().pull_dataframe()
+
+    # Mark rows that are in the model
+    model_ids = set(model_df[id_column])
+    full_df["in_model"] = full_df[id_column].isin(model_ids)
 
+    # Create and return the NoiseModel
+    return NoiseModel(full_df, id_column, features, target)
 
-def uq_model(model: "Model", uq_model_name: str, train_all_data: bool = False) -> "Model":
-    """Create a Uncertainty Quantification (UQ) model based on the given model
+
+def cleanlab_model_local(model: "Model"):
+    """Create a CleanlabModels instance for detecting data quality issues in a Model's training data.
 
     Args:
-        model (Model): The model to create the UQ model from
-        uq_model_name (str): The name of the UQ model to create
-        train_all_data (bool, optional): Whether to train the UQ model on all data (default: False)
+        model (Model): The Model used to create the cleanlab models
 
     Returns:
-        Model: The UQ model
+        CleanlabModels: Factory providing access to CleanLearning and Datalab models.
+            - clean_learning(): CleanLearning model with enhanced get_label_issues()
+            - datalab(): Datalab instance with report(), get_issues()
+    """
+    from workbench.algorithms.models.cleanlab_model import create_cleanlab_model  # noqa: F401 (avoid circular import)
+    from workbench.api import Model, FeatureSet  # noqa: F401 (avoid circular import)
+
+    # Get Feature and Target Columns from the existing given Model
+    features = model.features()
+    target = model.target()
+    model_type = model.model_type
+
+    # Backtrack our FeatureSet to get the ID column
+    fs = FeatureSet(model.get_input())
+    id_column = fs.id_column
+
+    # Get the full FeatureSet data
+    full_df = fs.pull_dataframe()
+
+    # Create and return the CleanLearning model
+    return create_cleanlab_model(full_df, id_column, features, target, model_type=model_type)
+
+
+def published_proximity_model(model: "Model", prox_model_name: str, include_all_columns: bool = False) -> "Model":
+    """Create a published proximity model based on the given model
+
+    Args:
+        model (Model): The model to create the proximity model from
+        prox_model_name (str): The name of the proximity model to create
+        include_all_columns (bool): Include all DataFrame columns in results (default: False)
+    Returns:
+        Model: The proximity model
     """
     from workbench.api import Model, ModelType, FeatureSet  # noqa: F401 (avoid circular import)
 
+    # Get the custom script path for the proximity model
+    script_path = get_custom_script_path("proximity", "feature_space_proximity.template")
+
     # Get Feature and Target Columns from the existing given Model
     features = model.features()
     target = model.target()
 
     # Create the Proximity Model from our FeatureSet
     fs = FeatureSet(model.get_input())
-    uq_model = fs.to_model(
-        name=uq_model_name,
-        model_type=ModelType.UQ_REGRESSOR,
+    prox_model = fs.to_model(
+        name=prox_model_name,
+        model_type=ModelType.PROXIMITY,
         feature_list=features,
         target_column=target,
-        description=f"UQ Model for {model.name}",
-        tags=["uq", model.name],
-        train_all_data=train_all_data,
-        custom_args={"id_column": fs.id_column, "track_columns": [target]},
+        description=f"Proximity Model for {model.name}",
+        tags=["proximity", model.name],
+        custom_script=script_path,
+        custom_args={"include_all_columns": include_all_columns},
     )
-    return uq_model
+    return prox_model
 
 
 def safe_extract_tarfile(tar_path: str, extract_path: str) -> None:
@@ -239,6 +320,63 @@ def load_category_mappings_from_s3(model_artifact_uri: str) -> Optional[dict]:
     return category_mappings
 
 
+def load_hyperparameters_from_s3(model_artifact_uri: str) -> Optional[dict]:
+    """
+    Download and extract hyperparameters from a model artifact in S3.
+
+    Args:
+        model_artifact_uri (str): S3 URI of the model artifact (model.tar.gz).
+
+    Returns:
+        dict: The loaded hyperparameters or None if not found.
+    """
+    hyperparameters = None
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Download model artifact
+        local_tar_path = os.path.join(tmpdir, "model.tar.gz")
+        wr.s3.download(path=model_artifact_uri, local_file=local_tar_path)
+
+        # Extract tarball
+        safe_extract_tarfile(local_tar_path, tmpdir)
+
+        # Look for hyperparameters in base directory only
+        hyperparameters_path = os.path.join(tmpdir, "hyperparameters.json")
+
+        if os.path.exists(hyperparameters_path):
+            try:
+                with open(hyperparameters_path, "r") as f:
+                    hyperparameters = json.load(f)
+                log.info(f"Loaded hyperparameters from {hyperparameters_path}")
+            except Exception as e:
+                log.warning(f"Failed to load hyperparameters from {hyperparameters_path}: {e}")
+
+    return hyperparameters
+
+
+def get_model_hyperparameters(workbench_model: Any) -> Optional[dict]:
+    """Get the hyperparameters used to train a Workbench model.
+
+    This retrieves the hyperparameters.json file from the model artifacts
+    that was saved during model training.
+
+    Args:
+        workbench_model: Workbench model object
+
+    Returns:
+        dict: The hyperparameters used during training, or None if not found
+    """
+    # Get the model artifact URI
+    model_artifact_uri = workbench_model.model_data_url()
+
+    if model_artifact_uri is None:
+        log.warning(f"No model artifact found for {workbench_model.uuid}")
+        return None
+
+    log.info(f"Loading hyperparameters from {model_artifact_uri}")
+    return load_hyperparameters_from_s3(model_artifact_uri)
+
+
 def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     """
     Evaluate uncertainty quantification model with essential metrics.
@@ -259,6 +397,13 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     if "prediction" not in df.columns:
         raise ValueError("Prediction column 'prediction' not found in DataFrame.")
 
+    # Drop rows with NaN predictions (e.g., from models that can't handle missing features)
+    n_total = len(df)
+    df = df.dropna(subset=["prediction", target_col])
+    n_valid = len(df)
+    if n_valid < n_total:
+        log.info(f"UQ metrics: dropped {n_total - n_valid} rows with NaN predictions")
+
     # --- Coverage and Interval Width ---
     if "q_025" in df.columns and "q_975" in df.columns:
         lower_95, upper_95 = df["q_025"], df["q_975"]
@@ -350,7 +495,7 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
 
 if __name__ == "__main__":
     """Exercise the Model Utilities"""
-    from workbench.api import Model, Endpoint
+    from workbench.api import Model
 
     # Get the instance information
     print(model_instance_info())
@@ -365,18 +510,11 @@ if __name__ == "__main__":
     # Get the custom script path
     print(get_custom_script_path("chem_info", "molecular_descriptors.py"))
 
-    # Test the proximity model
+    # Test loading hyperparameters
     m = Model("aqsol-regression")
+    hyperparams = get_model_hyperparameters(m)
+    print(hyperparams)
+
+    # Test the proximity model
     # prox_model = proximity_model(m, "aqsol-prox")
     # print(prox_model)#
-
-    # Test the UQ model
-    # uq_model_instance = uq_model(m, "aqsol-uq")
-    # print(uq_model_instance)
-    # uq_model_instance.to_endpoint()
-
-    # Test the uq_metrics function
-    end = Endpoint("aqsol-uq")
-    df = end.auto_inference(capture=True)
-    results = uq_metrics(df, target_col="solubility")
-    print(results)