workbench 0.8.168__py3-none-any.whl → 0.8.193__py3-none-any.whl

workbench/core/artifacts/feature_set_core.py

@@ -17,7 +17,7 @@ from workbench.core.artifacts.artifact import Artifact
 from workbench.core.artifacts.data_source_factory import DataSourceFactory
 from workbench.core.artifacts.athena_source import AthenaSource
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, List, Union
 
 from workbench.utils.aws_utils import aws_throttle
 
@@ -194,24 +194,24 @@ class FeatureSetCore(Artifact):
 
         return View(self, view_name)
 
-    def set_display_columns(self, diplay_columns: list[str]):
+    def set_display_columns(self, display_columns: list[str]):
         """Set the display columns for this Data Source
 
         Args:
-            diplay_columns (list[str]): The display columns for this Data Source
+            display_columns (list[str]): The display columns for this Data Source
         """
         # Check mismatch of display columns to computation columns
         c_view = self.view("computation")
         computation_columns = c_view.columns
-        mismatch_columns = [col for col in diplay_columns if col not in computation_columns]
+        mismatch_columns = [col for col in display_columns if col not in computation_columns]
         if mismatch_columns:
             self.log.monitor(f"Display View/Computation mismatch: {mismatch_columns}")
 
-        self.log.important(f"Setting Display Columns...{diplay_columns}")
+        self.log.important(f"Setting Display Columns...{display_columns}")
         from workbench.core.views import DisplayView
 
         # Create a NEW display view
-        DisplayView.create(self, source_table=c_view.table, column_list=diplay_columns)
+        DisplayView.create(self, source_table=c_view.table, column_list=display_columns)
 
     def set_computation_columns(self, computation_columns: list[str], reset_display: bool = True):
         """Set the computation columns for this FeatureSet
@@ -509,6 +509,48 @@ class FeatureSetCore(Artifact):
         ].tolist()
         return hold_out_ids
 
+    def set_training_filter(self, filter_expression: Optional[str] = None):
+        """Set a filter expression for the training view for this FeatureSet
+
+        Args:
+            filter_expression (Optional[str]): A SQL filter expression (e.g., "age > 25 AND status = 'active'")
+                If None or empty string, will reset to training view with no filter
+                (default: None)
+        """
+        from workbench.core.views import TrainingView
+
+        # Grab the existing holdout ids
+        holdout_ids = self.get_training_holdouts()
+
+        # Create a NEW training view
+        self.log.important(f"Setting Training Filter: {filter_expression}")
+        TrainingView.create(
+            self, id_column=self.id_column, holdout_ids=holdout_ids, filter_expression=filter_expression
+        )
+
+    def exclude_ids_from_training(self, ids: List[Union[str, int]], column_name: Optional[str] = None):
+        """Exclude a list of IDs from the training view
+
+        Args:
+            ids (List[Union[str, int]],): List of IDs to exclude from training
+            column_name (Optional[str]): Column name to filter on.
+                If None, uses self.id_column (default: None)
+        """
+        # Use the default id_column if not specified
+        column = column_name or self.id_column
+
+        # Handle empty list case
+        if not ids:
+            self.log.warning("No IDs provided to exclude")
+            return
+
+        # Build the filter expression with proper SQL quoting
+        quoted_ids = ", ".join([repr(id) for id in ids])
+        filter_expression = f"{column} NOT IN ({quoted_ids})"
+
+        # Apply the filter
+        self.set_training_filter(filter_expression)
+
     @classmethod
     def delete_views(cls, table: str, database: str):
         """Delete any views associated with this FeatureSet
@@ -707,7 +749,7 @@ if __name__ == "__main__":
 
     # Test getting the holdout ids
     print("Getting the hold out ids...")
-    holdout_ids = my_features.get_training_holdouts("id")
+    holdout_ids = my_features.get_training_holdouts()
     print(f"Holdout IDs: {holdout_ids}")
 
     # Get a sample of the data
@@ -729,16 +771,33 @@ if __name__ == "__main__":
     table = my_features.view("training").table
     df = my_features.query(f'SELECT id, name FROM "{table}"')
     my_holdout_ids = [id for id in df["id"] if id < 20]
-    my_features.set_training_holdouts("id", my_holdout_ids)
-
-    # Test the hold out set functionality with strings
-    print("Setting hold out ids (strings)...")
-    my_holdout_ids = [name for name in df["name"] if int(name.split(" ")[1]) > 80]
-    my_features.set_training_holdouts("name", my_holdout_ids)
+    my_features.set_training_holdouts(my_holdout_ids)
 
     # Get the training data
     print("Getting the training data...")
     training_data = my_features.get_training_data()
+    print(f"Training Data: {training_data.shape}")
+
+    # Test the filter expression functionality
+    print("Setting a filter expression...")
+    my_features.set_training_filter("id < 50 AND height > 65.0")
+    training_data = my_features.get_training_data()
+    print(f"Training Data: {training_data.shape}")
+    print(training_data)
+
+    # Remove training filter
+    print("Removing the filter expression...")
+    my_features.set_training_filter(None)
+    training_data = my_features.get_training_data()
+    print(f"Training Data: {training_data.shape}")
+    print(training_data)
+
+    # Test excluding ids from training
+    print("Excluding ids from training...")
+    my_features.exclude_ids_from_training([1, 2, 3, 4, 5])
+    training_data = my_features.get_training_data()
+    print(f"Training Data: {training_data.shape}")
+    print(training_data)
 
     # Now delete the AWS artifacts associated with this Feature Set
     # print("Deleting Workbench Feature Set...")

workbench/core/artifacts/model_core.py

@@ -21,6 +21,7 @@ from workbench.utils.aws_utils import newest_path, pull_s3_data
 from workbench.utils.s3_utils import compute_s3_object_hash
 from workbench.utils.shap_utils import shap_values_data, shap_feature_importance
 from workbench.utils.deprecated_utils import deprecated
+from workbench.utils.model_utils import proximity_model
 
 
 class ModelType(Enum):
@@ -40,58 +41,39 @@ class ModelType(Enum):
 class ModelImages:
     """Class for retrieving workbench inference images"""
 
-    image_uris = {
-        # US East 1 images
-        ("us-east-1", "xgb_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-training:0.1"
-        ),
-        ("us-east-1", "xgb_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1"
-        ),
-        ("us-east-1", "pytorch_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-pytorch-training:0.1"
-        ),
-        ("us-east-1", "pytorch_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-pytorch-inference:0.1"
-        ),
-        # US West 2 images
-        ("us-west-2", "xgb_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-training:0.1"
-        ),
-        ("us-west-2", "xgb_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1"
-        ),
-        ("us-west-2", "pytorch_training", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-pytorch-training:0.1"
-        ),
-        ("us-west-2", "pytorch_inference", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-pytorch-inference:0.1"
-        ),
-        # ARM64 images
-        ("us-east-1", "xgb_inference", "0.1", "arm64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1-arm64"
-        ),
-        ("us-west-2", "xgb_inference", "0.1", "arm64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-sklearn-xgb-inference:0.1-arm64"
-        ),
-        # Meta Endpoint inference images
-        ("us-east-1", "meta-endpoint", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-east-1.amazonaws.com/aws-ml-images/py312-meta-endpoint:0.1"
-        ),
-        ("us-west-2", "meta-endpoint", "0.1", "x86_64"): (
-            "507740646243.dkr.ecr.us-west-2.amazonaws.com/aws-ml-images/py312-meta-endpoint:0.1"
-        ),
+    # Account ID
+    ACCOUNT_ID = "507740646243"
+
+    # Image name mappings
+    IMAGE_NAMES = {
+        "training": "py312-general-ml-training",
+        "inference": "py312-general-ml-inference",
+        "pytorch_training": "py312-pytorch-training",
+        "pytorch_inference": "py312-pytorch-inference",
+        "meta-endpoint": "py312-meta-endpoint",
     }
 
     @classmethod
     def get_image_uri(cls, region, image_type, version="0.1", architecture="x86_64"):
-        key = (region, image_type, version, architecture)
-        if key in cls.image_uris:
-            return cls.image_uris[key]
-        else:
-            raise ValueError(
-                f"No matching image found for region: {region}, image_type: {image_type}, version: {version}"
-            )
+        """
+        Dynamically construct ECR image URI.
+
+        Args:
+            region: AWS region (e.g., 'us-east-1', 'us-west-2')
+            image_type: Type of image (e.g., 'training', 'inference', 'pytorch_training')
+            version: Image version (e.g., '0.1', '0.2')
+            architecture: CPU architecture (default: 'x86_64', currently unused but kept for compatibility)
+
+        Returns:
+            ECR image URI string
+        """
+        if image_type not in cls.IMAGE_NAMES:
+            raise ValueError(f"Unknown image_type: {image_type}. Valid types: {list(cls.IMAGE_NAMES.keys())}")
+
+        image_name = cls.IMAGE_NAMES[image_type]
+        uri = f"{cls.ACCOUNT_ID}.dkr.ecr.{region}.amazonaws.com/aws-ml-images/{image_name}:{version}"
+
+        return uri
 
 
 class ModelCore(Artifact):
@@ -597,6 +579,24 @@ class ModelCore(Artifact):
         # Return the details
         return details
 
+    # Training View for this model
+    def training_view(self):
+        """Get the training view for this model"""
+        from workbench.core.artifacts.feature_set_core import FeatureSetCore
+        from workbench.core.views import View
+
+        # Grab our FeatureSet
+        fs = FeatureSetCore(self.get_input())
+
+        # See if we have a training view for this model
+        my_model_training_view = f"{self.name.replace('-', '_')}_training"
+        view = View(fs, my_model_training_view, auto_create_view=False)
+        if view.exists():
+            return view
+        else:
+            self.log.important(f"No specific training view {my_model_training_view}, returning default training view")
+            return fs.view("training")
+
     # Pipeline for this model
     def get_pipeline(self) -> str:
         """Get the pipeline for this model"""
@@ -867,6 +867,14 @@ class ModelCore(Artifact):
                 shap_data[key] = self.df_store.get(df_location)
             return shap_data or None
 
+    def cross_folds(self) -> dict:
+        """Retrieve the cross-fold inference results(only works for XGBoost models)
+
+        Returns:
+            dict: Dictionary with the cross-fold inference results
+        """
+        return self.param_store.get(f"/workbench/models/{self.name}/inference/cross_fold")
+
     def supported_inference_instances(self) -> Optional[list]:
         """Retrieve the supported endpoint inference instance types
 
@@ -879,10 +887,24 @@ class ModelCore(Artifact):
         except (KeyError, IndexError, TypeError):
             return None
 
+    def publish_prox_model(self, prox_model_name: str = None, track_columns: list = None):
+        """Create and publish a Proximity Model for this Model
+
+        Args:
+            prox_model_name (str, optional): Name of the Proximity Model (if not specified, a name will be generated)
+            track_columns (list, optional): List of columns to track in the Proximity Model.
+
+        Returns:
+            Model: The published Proximity Model
+        """
+        if prox_model_name is None:
+            prox_model_name = self.model_name + "-prox"
+        return proximity_model(self, prox_model_name, track_columns=track_columns)
+
     def delete(self):
         """Delete the Model Packages and the Model Group"""
         if not self.exists():
-            self.log.warning(f"Trying to delete an Model that doesn't exist: {self.name}")
+            self.log.warning(f"Trying to delete a Model that doesn't exist: {self.name}")
 
         # Call the Class Method to delete the Model Group
         ModelCore.managed_delete(model_group_name=self.name)