workbench 0.8.177__py3-none-any.whl → 0.8.227__py3-none-any.whl

workbench/cached/cached_model.py

@@ -4,7 +4,7 @@ from typing import Union
 import pandas as pd
 
 # Workbench Imports
-from workbench.core.artifacts.model_core import ModelCore
+from workbench.core.artifacts.model_core import ModelCore, ModelType
 from workbench.core.artifacts.cached_artifact_mixin import CachedArtifactMixin
 
 
@@ -72,11 +72,11 @@ class CachedModel(CachedArtifactMixin, ModelCore):
         return super().list_inference_runs()
 
     @CachedArtifactMixin.cache_result
-    def get_inference_metrics(self, capture_name: str = "latest") -> Union[pd.DataFrame, None]:
+    def get_inference_metrics(self, capture_name: str = "auto") -> Union[pd.DataFrame, None]:
         """Retrieve the captured prediction results for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name (default: latest)
+            capture_name (str, optional): Specific capture_name (default: auto)
 
         Returns:
             pd.DataFrame: DataFrame of the Captured Metrics (might be None)
@@ -84,28 +84,66 @@ class CachedModel(CachedArtifactMixin, ModelCore):
         return super().get_inference_metrics(capture_name=capture_name)
 
     @CachedArtifactMixin.cache_result
-    def get_inference_predictions(self, capture_name: str = "auto_inference") -> Union[pd.DataFrame, None]:
+    def get_inference_predictions(
+        self, capture_name: str = "auto_inference", limit: int = 1000
+    ) -> Union[pd.DataFrame, None]:
         """Retrieve the captured prediction results for this model
 
         Args:
-            capture_name (str, optional): Specific capture_name (default: training_holdout)
+            capture_name (str, optional): Specific capture_name (default: auto_inference)
+            limit (int, optional): Maximum rows to return (default: 1000)
 
         Returns:
             pd.DataFrame: DataFrame of the Captured Predictions (might be None)
         """
-        # Note: This method can generate larger dataframes, so we'll sample if needed
         df = super().get_inference_predictions(capture_name=capture_name)
-        if df is not None and len(df) > 5000:
-            self.log.warning(f"{self.name}:{capture_name} Sampling Inference Predictions to 5000 rows")
-            return df.sample(5000)
+        if df is None:
+            return None
+
+        # Compute residual and do smart sampling based on model type
+        is_regressor = self.model_type in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.ENSEMBLE_REGRESSOR]
+        is_classifier = self.model_type == ModelType.CLASSIFIER
+
+        if is_regressor:
+            target = self.target()
+            if target and "prediction" in df.columns and target in df.columns:
+                df["residual"] = abs(df["prediction"] - df[target])
+
+        elif is_classifier:
+            target = self.target()
+            class_labels = self.class_labels()
+            if target and "prediction" in df.columns and target in df.columns and class_labels:
+                # Create a mapping from label to ordinal index
+                label_to_idx = {label: idx for idx, label in enumerate(class_labels)}
+                # Compute residual as distance between predicted and actual class
+                df["residual"] = abs(
+                    df["prediction"].map(label_to_idx).fillna(-1) - df[target].map(label_to_idx).fillna(-1)
+                )
+
+        # Smart sampling: half high-residual rows, half random from the rest
+        if "residual" in df.columns and len(df) > limit:
+            half_limit = limit // 2
+            self.log.warning(
+                f"{self.name}:{capture_name} Sampling {limit} rows (top {half_limit} residuals + {half_limit} random)"
+            )
+            top_residuals = df.nlargest(half_limit, "residual")
+            remaining = df.drop(top_residuals.index)
+            random_sample = remaining.sample(min(half_limit, len(remaining)))
+            return pd.concat([top_residuals, random_sample]).reset_index(drop=True)
+
+        # Fallback: just limit rows if no residual computed
+        if len(df) > limit:
+            self.log.warning(f"{self.name}:{capture_name} Sampling to {limit} rows")
+            return df.sample(limit)
+
         return df
 
     @CachedArtifactMixin.cache_result
-    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 the model
 
         Args:
-            capture_name (str, optional): Specific capture_name (default: latest)
+            capture_name (str, optional): Specific capture_name (default: auto)
 
         Returns:
             pd.DataFrame: DataFrame of the Confusion Matrix (might be None)

workbench/core/artifacts/__init__.py

@@ -15,7 +15,16 @@ from .artifact import Artifact
 from .athena_source import AthenaSource
 from .data_source_abstract import DataSourceAbstract
 from .feature_set_core import FeatureSetCore
-from .model_core import ModelCore, ModelType
+from .model_core import ModelCore, ModelType, ModelFramework
 from .endpoint_core import EndpointCore
 
-__all__ = ["Artifact", "AthenaSource", "DataSourceAbstract", "FeatureSetCore", "ModelCore", "ModelType", "EndpointCore"]
+__all__ = [
+    "Artifact",
+    "AthenaSource",
+    "DataSourceAbstract",
+    "FeatureSetCore",
+    "ModelCore",
+    "ModelType",
+    "ModelFramework",
+    "EndpointCore",
+]

workbench/core/artifacts/artifact.py

@@ -8,8 +8,8 @@ from typing import Union
 
 # Workbench Imports
 from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
-from workbench.core.cloud_platform.aws.aws_parameter_store import AWSParameterStore as ParameterStore
-from workbench.core.cloud_platform.aws.aws_df_store import AWSDFStore as DFStore
+from workbench.core.artifacts.parameter_store_core import ParameterStoreCore
+from workbench.core.artifacts.df_store_core import DFStoreCore
 from workbench.utils.aws_utils import dict_to_aws_tags
 from workbench.utils.config_manager import ConfigManager, FatalConfigError
 from workbench.core.cloud_platform.cloud_meta import CloudMeta
@@ -48,11 +48,11 @@ class Artifact(ABC):
     tag_delimiter = "::"
 
     # Grab our Dataframe Cache Storage
-    df_cache = DFStore(path_prefix="/workbench/dataframe_cache")
+    df_cache = DFStoreCore(path_prefix="/workbench/dataframe_cache")
 
     # Artifact may want to use the Parameter Store or Dataframe Store
-    param_store = ParameterStore()
-    df_store = DFStore()
+    param_store = ParameterStoreCore()
+    df_store = DFStoreCore()
 
     def __init__(self, name: str, use_cached_meta: bool = False):
         """Initialize the Artifact Base Class

workbench/core/artifacts/df_store_core.py

@@ -0,0 +1,114 @@
+"""DFStoreCore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy"""
+
+import logging
+from typing import Union
+
+# Workbench Imports
+from workbench.utils.config_manager import ConfigManager
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Bridges Import
+from workbench_bridges.api import DFStore as BridgesDFStore
+
+
+class DFStoreCore(BridgesDFStore):
+    """DFStoreCore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy
+
+        Common Usage:
+    ```python
+            df_store = DFStoreCore()
+
+            # List Data
+            df_store.list()
+
+            # Add DataFrame
+            df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+            df_store.upsert("/test/my_data", df)
+
+            # Retrieve DataFrame
+            df = df_store.get("/test/my_data")
+            print(df)
+
+            # Delete Data
+            df_store.delete("/test/my_data")
+    ```
+    """
+
+    def __init__(self, path_prefix: Union[str, None] = None):
+        """DFStoreCore Init Method
+
+        Args:
+            path_prefix (Union[str, None], optional): Add a path prefix to storage locations (Defaults to None)
+        """
+        # Get config from workbench's systems
+        bucket = ConfigManager().get_config("WORKBENCH_BUCKET")
+        session = AWSAccountClamp().boto3_session
+
+        # Initialize parent with workbench config
+        super().__init__(path_prefix=path_prefix, s3_bucket=bucket, boto3_session=session)
+        self.log = logging.getLogger("workbench")
+
+
+if __name__ == "__main__":
+    """Exercise the DFStoreCore Class"""
+    import time
+    import pandas as pd
+
+    # Create a DFStoreCore manager
+    df_store = DFStoreCore()
+
+    # Details of the Dataframe Store
+    print("Detailed Data...")
+    print(df_store.details())
+
+    # Add a new DataFrame
+    my_df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+    df_store.upsert("/testing/test_data", my_df)
+
+    # Get the DataFrame
+    print(f"Getting data 'test_data':\n{df_store.get('/testing/test_data')}")
+
+    # Now let's test adding a Series
+    series = pd.Series([1, 2, 3, 4], name="Series")
+    df_store.upsert("/testing/test_series", series)
+    print(f"Getting data 'test_series':\n{df_store.get('/testing/test_series')}")
+
+    # Summary of the data
+    print("Summary Data...")
+    print(df_store.summary())
+
+    # Repr of the DFStoreCore object
+    print("DFStoreCore Object:")
+    print(df_store)
+
+    # Check if the data exists
+    print("Check if data exists...")
+    print(df_store.check("/testing/test_data"))
+    print(df_store.check("/testing/test_series"))
+
+    # Time the check
+    start_time = time.time()
+    print(df_store.check("/testing/test_data"))
+    print("--- Check %s seconds ---" % (time.time() - start_time))
+
+    # Now delete the test data
+    df_store.delete("/testing/test_data")
+    df_store.delete("/testing/test_series")
+
+    # Check if the data exists
+    print("Check if data exists...")
+    print(df_store.check("/testing/test_data"))
+    print(df_store.check("/testing/test_series"))
+
+    # Add a bunch of dataframes and then test recursive delete
+    for i in range(10):
+        df_store.upsert(f"/testing/data_{i}", pd.DataFrame({"A": [1, 2], "B": [3, 4]}))
+    print("Before Recursive Delete:")
+    print(df_store.summary())
+    df_store.delete_recursive("/testing")
+    print("After Recursive Delete:")
+    print(df_store.summary())
+
+    # Get a non-existent DataFrame
+    print("Getting non-existent data...")
+    print(df_store.get("/testing/no_where"))