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

workbench/api/feature_set.py

@@ -58,10 +58,7 @@ class FeatureSet(FeatureSetCore):
             include_aws_columns (bool): Include the AWS columns in the DataFrame (default: False)
 
         Returns:
-            pd.DataFrame: A DataFrame of ALL the data from this FeatureSet
-
-        Note:
-            Obviously this is not recommended for large datasets :)
+            pd.DataFrame: A DataFrame of all the data from this FeatureSet up to the limit
         """
 
         # Get the table associated with the data
@@ -83,7 +80,7 @@ class FeatureSet(FeatureSetCore):
         tags: list = None,
         description: str = None,
         feature_list: list = None,
-        target_column: str = None,
+        target_column: Union[str, list[str]] = None,
         model_class: str = None,
         model_import_str: str = None,
         custom_script: Union[str, Path] = None,
@@ -103,7 +100,7 @@ class FeatureSet(FeatureSetCore):
             tags (list, optional): Set the tags for the model.  If not given tags will be generated.
             description (str, optional): Set the description for the model. If not give a description is generated.
             feature_list (list, optional): Set the feature list for the model. If not given a feature list is generated.
-            target_column (str, optional): The target column for the model (use None for unsupervised model)
+            target_column (str or list[str], optional): Target column(s) for the model (use None for unsupervised model)
             model_class (str, optional): Model class to use (e.g. "KMeans", default: None)
             model_import_str (str, optional): The import for the model (e.g. "from sklearn.cluster import KMeans")
             custom_script (str, optional): The custom script to use for the model (default: None)
@@ -131,7 +128,7 @@ class FeatureSet(FeatureSetCore):
         tags = [name] if tags is None else tags
 
         # If the model framework is PyTorch or ChemProp, ensure we set the training and inference images
-        if model_framework in (ModelFramework.PYTORCH_TABULAR, ModelFramework.CHEMPROP):
+        if model_framework in (ModelFramework.PYTORCH, ModelFramework.CHEMPROP):
             training_image = "pytorch_training"
             inference_image = "pytorch_inference"
 
@@ -157,6 +154,93 @@ class FeatureSet(FeatureSetCore):
         # Return the Model
         return Model(name)
 
+    def prox_model(
+        self, target: str, features: list, include_all_columns: bool = False
+    ) -> "FeatureSpaceProximity":  # noqa: F821
+        """Create a local FeatureSpaceProximity Model for this FeatureSet
+
+        Args:
+           target (str): The target column name
+           features (list): The list of feature column names
+           include_all_columns (bool): Include all DataFrame columns in results (default: False)
+
+        Returns:
+           FeatureSpaceProximity: A local FeatureSpaceProximity Model
+        """
+        from workbench.algorithms.dataframe.feature_space_proximity import FeatureSpaceProximity  # noqa: F401
+
+        # Create the Proximity Model from the full FeatureSet dataframe
+        full_df = self.pull_dataframe()
+
+        # Create and return the FeatureSpaceProximity Model
+        return FeatureSpaceProximity(
+            full_df, id_column=self.id_column, features=features, target=target, include_all_columns=include_all_columns
+        )
+
+    def fp_prox_model(
+        self,
+        target: str,
+        fingerprint_column: str = None,
+        include_all_columns: bool = False,
+        radius: int = 2,
+        n_bits: int = 1024,
+        counts: bool = False,
+    ) -> "FingerprintProximity":  # noqa: F821
+        """Create a local FingerprintProximity Model for this FeatureSet
+
+        Args:
+           target (str): The target column name
+           fingerprint_column (str): Column containing fingerprints. If None, uses existing 'fingerprint'
+                                     column or computes from SMILES column.
+           include_all_columns (bool): Include all DataFrame columns in results (default: False)
+           radius (int): Radius for Morgan fingerprint computation (default: 2)
+           n_bits (int): Number of bits for fingerprint (default: 1024)
+           counts (bool): Whether to use count simulation (default: False)
+
+        Returns:
+           FingerprintProximity: A local FingerprintProximity Model
+        """
+        from workbench.algorithms.dataframe.fingerprint_proximity import FingerprintProximity  # noqa: F401
+
+        # Create the Proximity Model from the full FeatureSet dataframe
+        full_df = self.pull_dataframe()
+
+        # Create and return the FingerprintProximity Model
+        return FingerprintProximity(
+            full_df,
+            id_column=self.id_column,
+            fingerprint_column=fingerprint_column,
+            target=target,
+            include_all_columns=include_all_columns,
+            radius=radius,
+            n_bits=n_bits,
+        )
+
+    def cleanlab_model(
+        self,
+        target: str,
+        features: list,
+        model_type: ModelType = ModelType.REGRESSOR,
+    ) -> "CleanLearning":  # noqa: F821
+        """Create a CleanLearning model for detecting label issues in this FeatureSet
+
+        Args:
+           target (str): The target column name
+           features (list): The list of feature column names
+           model_type (ModelType): The model type (REGRESSOR or CLASSIFIER). Defaults to REGRESSOR.
+
+        Returns:
+           CleanLearning: A fitted cleanlab model. Use get_label_issues() to get
+           a DataFrame with id_column, label_quality, predicted_label, given_label, is_label_issue.
+        """
+        from workbench.algorithms.models.cleanlab_model import create_cleanlab_model  # noqa: F401
+
+        # Get the full FeatureSet dataframe
+        full_df = self.pull_dataframe()
+
+        # Create and return the CleanLearning model
+        return create_cleanlab_model(full_df, self.id_column, features, target, model_type=model_type)
+
 
 if __name__ == "__main__":
     """Exercise the FeatureSet Class"""
@@ -167,5 +251,24 @@ if __name__ == "__main__":
     pprint(my_features.summary())
     pprint(my_features.details())
 
+    # Pull the full DataFrame
+    df = my_features.pull_dataframe()
+    print(df.head())
+
+    # Create a Proximity Model from the FeatureSet
+    features = ["height", "weight", "age", "iq_score", "likes_dogs", "food"]
+    my_prox = my_features.prox_model(target="salary", features=features)
+    neighbors = my_prox.neighbors(42)
+    print("Neighbors for ID 42:")
+    print(neighbors)
+
     # Create a Model from the FeatureSet
-    my_model = my_features.to_model(name="test-model", model_type=ModelType.REGRESSOR, target_column="iq_score")
+    """
+    my_model = my_features.to_model(
+        name="test-model",
+        model_type=ModelType.REGRESSOR,
+        target_column="salary",
+        feature_list=features
+    )
+    pprint(my_model.summary())
+    """

workbench/api/meta_model.py

@@ -0,0 +1,289 @@
+"""MetaModel: A Model that aggregates predictions from multiple child endpoints.
+
+MetaModels don't train on feature data - they combine predictions from existing
+endpoints using confidence-weighted voting. This provides ensemble benefits
+across different model frameworks (XGBoost, PyTorch, ChemProp, etc.).
+"""
+
+from pathlib import Path
+import time
+import logging
+
+from sagemaker.estimator import Estimator
+
+# Workbench Imports
+from workbench.api.model import Model
+from workbench.api.endpoint import Endpoint
+from workbench.core.artifacts.model_core import ModelCore, ModelType, ModelFramework, ModelImages
+from workbench.core.artifacts.artifact import Artifact
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+from workbench.model_scripts.script_generation import generate_model_script
+from workbench.utils.config_manager import ConfigManager
+from workbench.utils.model_utils import supported_instance_types
+
+# Set up logging
+log = logging.getLogger("workbench")
+
+
+class MetaModel(Model):
+    """MetaModel: A Model that aggregates predictions from child endpoints.
+
+    Common Usage:
+        ```python
+        # Create a meta model from existing endpoints
+        meta = MetaModel.create(
+            name="my-meta-model",
+            child_endpoints=["endpoint-1", "endpoint-2", "endpoint-3"],
+            target_column="target"
+        )
+
+        # Deploy like any other model
+        endpoint = meta.to_endpoint()
+        ```
+    """
+
+    @classmethod
+    def create(
+        cls,
+        name: str,
+        child_endpoints: list[str],
+        target_column: str,
+        description: str = None,
+        tags: list[str] = None,
+    ) -> "MetaModel":
+        """Create a new MetaModel from a list of child endpoints.
+
+        Args:
+            name: Name for the meta model
+            child_endpoints: List of endpoint names to aggregate
+            target_column: Name of the target column (for metadata)
+            description: Optional description for the model
+            tags: Optional list of tags
+
+        Returns:
+            MetaModel: The created meta model
+        """
+        Artifact.is_name_valid(name, delimiter="-", lower_case=False)
+
+        # Validate endpoints and get lineage info from primary endpoint
+        feature_list, feature_set_name, model_weights = cls._validate_and_get_lineage(child_endpoints)
+
+        # Delete existing model if it exists
+        log.important(f"Trying to delete existing model {name}...")
+        ModelCore.managed_delete(name)
+
+        # Run training and register model
+        aws_clamp = AWSAccountClamp()
+        estimator = cls._run_training(name, child_endpoints, target_column, model_weights, aws_clamp)
+        cls._register_model(name, child_endpoints, description, tags, estimator, aws_clamp)
+
+        # Set metadata and onboard
+        cls._set_metadata(name, target_column, feature_list, feature_set_name, child_endpoints)
+
+        log.important(f"MetaModel {name} created successfully!")
+        return cls(name)
+
+    @classmethod
+    def _validate_and_get_lineage(cls, child_endpoints: list[str]) -> tuple[list[str], str, dict[str, float]]:
+        """Validate child endpoints exist and get lineage info from primary endpoint.
+
+        Args:
+            child_endpoints: List of endpoint names
+
+        Returns:
+            tuple: (feature_list, feature_set_name, model_weights) from the primary endpoint's model
+        """
+        log.info("Verifying child endpoints and gathering model metrics...")
+        mae_scores = {}
+
+        for ep_name in child_endpoints:
+            ep = Endpoint(ep_name)
+            if not ep.exists():
+                raise ValueError(f"Child endpoint '{ep_name}' does not exist")
+
+            # Get model MAE from full_inference metrics
+            model = Model(ep.get_input())
+            metrics = model.get_inference_metrics("full_inference")
+            if metrics is not None and "mae" in metrics.columns:
+                mae = float(metrics["mae"].iloc[0])
+                mae_scores[ep_name] = mae
+                log.info(f"  {ep_name} -> {model.name}: MAE={mae:.4f}")
+            else:
+                log.warning(f"  {ep_name}: No full_inference metrics found, using default weight")
+                mae_scores[ep_name] = None
+
+        # Compute inverse-MAE weights (higher weight for lower MAE)
+        valid_mae = {k: v for k, v in mae_scores.items() if v is not None}
+        if valid_mae:
+            inv_mae = {k: 1.0 / v for k, v in valid_mae.items()}
+            total = sum(inv_mae.values())
+            model_weights = {k: v / total for k, v in inv_mae.items()}
+            # Fill in missing weights with equal share of remaining weight
+            missing = [k for k in mae_scores if mae_scores[k] is None]
+            if missing:
+                equal_weight = (1.0 - sum(model_weights.values())) / len(missing)
+                for k in missing:
+                    model_weights[k] = equal_weight
+        else:
+            # No metrics available, use equal weights
+            model_weights = {k: 1.0 / len(child_endpoints) for k in child_endpoints}
+            log.warning("No MAE metrics found, using equal weights")
+
+        log.info(f"Model weights: {model_weights}")
+
+        # Use first endpoint as primary - backtrack to get model and feature set
+        primary_endpoint = Endpoint(child_endpoints[0])
+        primary_model = Model(primary_endpoint.get_input())
+        feature_list = primary_model.features()
+        feature_set_name = primary_model.get_input()
+
+        log.info(
+            f"Primary endpoint: {child_endpoints[0]} -> Model: {primary_model.name} -> FeatureSet: {feature_set_name}"
+        )
+        return feature_list, feature_set_name, model_weights
+
+    @classmethod
+    def _run_training(
+        cls,
+        name: str,
+        child_endpoints: list[str],
+        target_column: str,
+        model_weights: dict[str, float],
+        aws_clamp: AWSAccountClamp,
+    ) -> Estimator:
+        """Run the minimal training job that saves the meta model config.
+
+        Args:
+            name: Model name
+            child_endpoints: List of endpoint names
+            target_column: Target column name
+            model_weights: Dict mapping endpoint name to weight
+            aws_clamp: AWS account clamp
+
+        Returns:
+            Estimator: The fitted estimator
+        """
+        sm_session = aws_clamp.sagemaker_session()
+        cm = ConfigManager()
+        workbench_bucket = cm.get_config("WORKBENCH_BUCKET")
+        models_s3_path = f"s3://{workbench_bucket}/models"
+
+        # Generate the model script from template
+        template_params = {
+            "model_type": ModelType.REGRESSOR,
+            "model_framework": ModelFramework.META,
+            "child_endpoints": child_endpoints,
+            "target_column": target_column,
+            "model_weights": model_weights,
+            "model_metrics_s3_path": f"{models_s3_path}/{name}/training",
+            "aws_region": sm_session.boto_region_name,
+        }
+        script_path = generate_model_script(template_params)
+
+        # Create estimator
+        training_image = ModelImages.get_image_uri(sm_session.boto_region_name, "meta_training")
+        log.info(f"Using Meta Training Image: {training_image}")
+        estimator = Estimator(
+            entry_point=Path(script_path).name,
+            source_dir=str(Path(script_path).parent),
+            role=aws_clamp.aws_session.get_workbench_execution_role_arn(),
+            instance_count=1,
+            instance_type="ml.m5.large",
+            sagemaker_session=sm_session,
+            image_uri=training_image,
+        )
+
+        # Run training (no input data needed - just saves config)
+        log.important(f"Creating MetaModel {name}...")
+        estimator.fit()
+
+        return estimator
+
+    @classmethod
+    def _register_model(
+        cls,
+        name: str,
+        child_endpoints: list[str],
+        description: str,
+        tags: list[str],
+        estimator: Estimator,
+        aws_clamp: AWSAccountClamp,
+    ):
+        """Create model group and register the model.
+
+        Args:
+            name: Model name
+            child_endpoints: List of endpoint names
+            description: Model description
+            tags: Model tags
+            estimator: Fitted estimator
+            aws_clamp: AWS account clamp
+        """
+        sm_session = aws_clamp.sagemaker_session()
+        model_description = description or f"Meta model aggregating: {', '.join(child_endpoints)}"
+
+        # Create model group
+        aws_clamp.sagemaker_client().create_model_package_group(
+            ModelPackageGroupName=name,
+            ModelPackageGroupDescription=model_description,
+            Tags=[{"Key": "workbench_tags", "Value": "::".join(tags or [name])}],
+        )
+
+        # Register the model with meta inference image
+        inference_image = ModelImages.get_image_uri(sm_session.boto_region_name, "meta_inference")
+        log.important(f"Registering model {name} with Inference Image {inference_image}...")
+        estimator.create_model(role=aws_clamp.aws_session.get_workbench_execution_role_arn()).register(
+            model_package_group_name=name,
+            image_uri=inference_image,
+            content_types=["text/csv"],
+            response_types=["text/csv"],
+            inference_instances=supported_instance_types("x86_64"),
+            transform_instances=["ml.m5.large", "ml.m5.xlarge"],
+            approval_status="Approved",
+            description=model_description,
+        )
+
+    @classmethod
+    def _set_metadata(
+        cls, name: str, target_column: str, feature_list: list[str], feature_set_name: str, child_endpoints: list[str]
+    ):
+        """Set model metadata and onboard.
+
+        Args:
+            name: Model name
+            target_column: Target column name
+            feature_list: List of feature names
+            feature_set_name: Name of the input FeatureSet
+            child_endpoints: List of child endpoint names
+        """
+        time.sleep(3)
+        output_model = ModelCore(name)
+        output_model._set_model_type(ModelType.UQ_REGRESSOR)
+        output_model._set_model_framework(ModelFramework.META)
+        output_model.set_input(feature_set_name, force=True)
+        output_model.upsert_workbench_meta({"workbench_model_target": target_column})
+        output_model.upsert_workbench_meta({"workbench_model_features": feature_list})
+        output_model.upsert_workbench_meta({"child_endpoints": child_endpoints})
+        output_model.onboard_with_args(ModelType.UQ_REGRESSOR, target_column, feature_list=feature_list)
+
+
+if __name__ == "__main__":
+    """Exercise the MetaModel Class"""
+
+    meta = MetaModel.create(
+        name="logd-meta",
+        child_endpoints=["logd-xgb", "logd-pytorch", "logd-chemprop"],
+        target_column="logd",
+        description="Meta model for LogD prediction",
+        tags=["meta", "logd", "ensemble"],
+    )
+    print(meta.summary())
+
+    # Create an endpoint for the meta model
+    end = meta.to_endpoint(tags=["meta", "logd"])
+    end.set_owner("BW")
+    end.auto_inference()
+
+    # Test loading an existing meta model
+    meta = MetaModel("logd-meta")
+    print(meta.details())

workbench/api/model.py

@@ -10,7 +10,12 @@ from workbench.core.artifacts.artifact import Artifact
 from workbench.core.artifacts.model_core import ModelCore, ModelType, ModelFramework  # noqa: F401
 from workbench.core.transforms.model_to_endpoint.model_to_endpoint import ModelToEndpoint
 from workbench.api.endpoint import Endpoint
-from workbench.utils.model_utils import proximity_model_local, uq_model
+from workbench.utils.model_utils import (
+    proximity_model_local,
+    fingerprint_prox_model_local,
+    noise_model_local,
+    cleanlab_model_local,
+)
 
 
 class Model(ModelCore):
@@ -83,27 +88,55 @@ class Model(ModelCore):
         end.set_owner(self.get_owner())
         return end
 
-    def prox_model(self):
+    def prox_model(self, include_all_columns: bool = False):
         """Create a local Proximity Model for this Model
 
+        Args:
+              include_all_columns (bool): Include all DataFrame columns in results (default: False)
+
         Returns:
-           Proximity: A local Proximity Model
+           FeatureSpaceProximity: A local FeatureSpaceProximity Model
         """
-        return proximity_model_local(self)
+        return proximity_model_local(self, include_all_columns=include_all_columns)
 
-    def uq_model(self, uq_model_name: str = None, train_all_data: bool = False) -> "Model":
-        """Create a Uncertainty Quantification Model for this Model
+    def fp_prox_model(
+        self,
+        include_all_columns: bool = False,
+        radius: int = 2,
+        n_bits: int = 1024,
+        counts: bool = False,
+    ):
+        """Create a local Fingerprint Proximity Model for this Model
 
         Args:
-            uq_model_name (str, optional): Name of the UQ Model (if not specified, a name will be generated)
-            train_all_data (bool, optional): Whether to train the UQ Model on all data (default: False)
+           include_all_columns (bool): Include all DataFrame columns in 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:
+           FingerprintProximity: A local FingerprintProximity Model
+        """
+        return fingerprint_prox_model_local(
+            self, include_all_columns=include_all_columns, radius=radius, n_bits=n_bits, counts=counts
+        )
+
+    def noise_model(self):
+        """Create a local Noise Model for this Model
+
+        Returns:
+           NoiseModel: A local Noise Model
+        """
+        return noise_model_local(self)
+
+    def cleanlab_model(self):
+        """Create a CleanLearning model for this Model's training data.
 
         Returns:
-            Model: The UQ Model
+           CleanLearning: A fitted cleanlab model. Use get_label_issues() to get
+           a DataFrame with id_column, label_quality, predicted_label, given_label, is_label_issue.
         """
-        if uq_model_name is None:
-            uq_model_name = self.model_name + "-uq"
-        return uq_model(self, uq_model_name, train_all_data=train_all_data)
+        return cleanlab_model_local(self)
 
 
 if __name__ == "__main__":

workbench/api/parameter_store.py

@@ -1,13 +1,10 @@
 """ParameterStore: Manages Workbench parameters in a Cloud Based Parameter Store."""
 
-from typing import Union
-import logging
-
 # Workbench Imports
-from workbench.core.cloud_platform.aws.aws_parameter_store import AWSParameterStore
+from workbench.core.artifacts.parameter_store_core import ParameterStoreCore
 
 
-class ParameterStore(AWSParameterStore):
+class ParameterStore(ParameterStoreCore):
     """ParameterStore: Manages Workbench parameters in a Cloud Based Parameter Store.
 
     Common Usage:
@@ -43,56 +40,10 @@ class ParameterStore(AWSParameterStore):
 
     def __init__(self):
         """ParameterStore Init Method"""
-        self.log = logging.getLogger("workbench")
 
-        # Initialize the SuperClass
+        # Initialize parent class
         super().__init__()
 
-    def list(self, prefix: str = None) -> list:
-        """List all parameters in the AWS Parameter Store, optionally filtering by a prefix.
-
-        Args:
-            prefix (str, optional): A prefix to filter the parameters by. Defaults to None.
-
-        Returns:
-            list: A list of parameter names and details.
-        """
-        return super().list(prefix=prefix)
-
-    def get(self, name: str, warn: bool = True, decrypt: bool = True) -> Union[str, list, dict, None]:
-        """Retrieve a parameter value from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to retrieve.
-            warn (bool): Whether to log a warning if the parameter is not found.
-            decrypt (bool): Whether to decrypt secure string parameters.
-
-        Returns:
-            Union[str, list, dict, None]: The value of the parameter or None if not found.
-        """
-        return super().get(name=name, warn=warn, decrypt=decrypt)
-
-    def upsert(self, name: str, value):
-        """Insert or update a parameter in the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter.
-            value (str | list | dict): The value of the parameter.
-        """
-        super().upsert(name=name, value=value)
-
-    def delete(self, name: str):
-        """Delete a parameter from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to delete.
-        """
-        super().delete(name=name)
-
-    def __repr__(self):
-        """Return a string representation of the ParameterStore object."""
-        return super().__repr__()
-
 
 if __name__ == "__main__":
     """Exercise the ParameterStore Class"""

workbench/cached/cached_model.py

@@ -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)
@@ -101,11 +101,11 @@ class CachedModel(CachedArtifactMixin, ModelCore):
         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/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