workbench 0.8.241__tar.gz → 0.8.246__tar.gz

{workbench-0.8.241 → workbench-0.8.246}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.241
+Version: 0.8.246
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License: MIT License

{workbench-0.8.241 → workbench-0.8.246}/glue_jobs/glue_mixed_case.py

@@ -28,7 +28,7 @@ glue_client.create_table(
         "Name": table_name,
         "StorageDescriptor": {
             "Columns": [{"Name": name, "Type": dtype} for name, dtype in original_column_dict.items()],
-            "Location": "s3://test-bucket/test_table/",
+            "Location": "s3://workbench-public-test-bucket/test_table/",
             "InputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat",
             "OutputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat",
             "SerdeInfo": {"SerializationLibrary": "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe"},

{workbench-0.8.241 → workbench-0.8.246}/model_docker_images/ml_pipelines/Dockerfile

@@ -29,7 +29,7 @@ COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/pytho
 COPY --from=builder /usr/local/bin /usr/local/bin
 
 # Install workbench only
-RUN pip install --no-cache-dir workbench==0.8.240
+RUN pip install --no-cache-dir workbench==0.8.245
 
 # Add the script runner
 COPY ml_pipeline_runner.py /app/ml_pipeline_runner.py

{workbench-0.8.241 → workbench-0.8.246}/pyproject.toml

@@ -78,6 +78,9 @@ packages = { find = { where = ["src"] } }
     "themes/**"
 ]
 
+[tool.setuptools.exclude-package-data]
+"workbench" = ["model_scripts/**/generated_*.py"]
+
 [project.scripts]
 workbench = "workbench.repl.workbench_shell:launch_shell"
 cloud_watch = "workbench.scripts.monitor_cloud_watch:main"

{workbench-0.8.241 → workbench-0.8.246}/scripts/admin/athena_ddl_mixed_case.py

@@ -22,7 +22,7 @@ athena_client = boto3_session.client("athena")
 # Define the database, table name, and S3 path
 database_name = "workbench"
 table_name = "test_table"
-s3_path = "s3://test-bucket/test/"
+s3_path = "s3://workbench-public-test-bucket/test/"
 
 # Define the original column names with mixed case
 original_column_dict = {"Id": "int", "Name": "string", "Age": "int"}

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/api/feature_set.py

@@ -85,9 +85,6 @@ class FeatureSet(FeatureSetCore):
         model_import_str: str = None,
         custom_script: Union[str, Path] = None,
         custom_args: dict = None,
-        training_image: str = "training",
-        inference_image: str = "inference",
-        inference_arch: str = "x86_64",
         **kwargs,
     ) -> Union[Model, None]:
         """Create a Model from the FeatureSet
@@ -104,9 +101,6 @@ class FeatureSet(FeatureSetCore):
             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)
-            training_image (str, optional): The training image to use (default: "training")
-            inference_image (str, optional): The inference image to use (default: "inference")
-            inference_arch (str, optional): The architecture to use for inference (default: "x86_64")
             kwargs (dict, optional): Additional keyword arguments to pass to the model
 
         Returns:
@@ -127,10 +121,15 @@ class FeatureSet(FeatureSetCore):
         # Create the Model Tags
         tags = [name] if tags is None else tags
 
-        # If the model framework is PyTorch or ChemProp, ensure we set the training and inference images
+        # Set training/inference images based on model framework
         if model_framework in (ModelFramework.PYTORCH, ModelFramework.CHEMPROP):
             training_image = "pytorch_training"
             inference_image = "pytorch_inference"
+            inference_arch = "x86_64"
+        else:
+            training_image = "training"
+            inference_image = "inference"
+            inference_arch = "x86_64"
 
         # Transform the FeatureSet into a Model
         features_to_model = FeaturesToModel(

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/core/artifacts/endpoint_core.py

@@ -1138,7 +1138,7 @@ class EndpointCore(Artifact):
             cls.log.error("Error deleting endpoint.")
             raise e
 
-        time.sleep(5)  # Final sleep for AWS to fully register deletions
+        time.sleep(10)  # Final sleep for AWS to fully register deletions
 
     @classmethod
     def delete_endpoint_models(cls, endpoint_name: str):

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/core/artifacts/feature_set_core.py

@@ -320,7 +320,7 @@ class FeatureSetCore(Artifact):
         """
 
         # Set up the S3 Query results path
-        date_time = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H:%M:%S")
+        date_time = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S")
         s3_output_path = self.feature_sets_s3_path + f"/{self.name}/datasets/all_{date_time}"
 
         # Make the query

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/core/transforms/features_to_model/features_to_model.py

@@ -247,9 +247,12 @@ class FeaturesToModel(Transform):
         # Create a Sagemaker Model with our script
         image = ModelImages.get_image_uri(self.sm_session.boto_region_name, self.training_image)
 
-        # Use GPU instance for ChemProp/PyTorch, CPU for others
-        if self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH]:
-            train_instance_type = "ml.g6.xlarge"  # NVIDIA L4 GPU, ~$0.80/hr
+        # Use user-specified instance or default based on framework
+        train_instance_type = kwargs.get("training_instance")
+        if train_instance_type:
+            self.log.important(f"Using user-specified instance {train_instance_type}")
+        elif self.model_framework in [ModelFramework.CHEMPROP, ModelFramework.PYTORCH]:
+            train_instance_type = "ml.g6.2xlarge"  # NVIDIA L4 GPU + 8 vCPUs for data loading
             self.log.important(f"Using GPU instance {train_instance_type} for {self.model_framework.value}")
         else:
             train_instance_type = "ml.m5.xlarge"

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/core/transforms/model_to_endpoint/model_to_endpoint.py

@@ -108,7 +108,6 @@ class ModelToEndpoint(Transform):
 
         # Is this a serverless deployment?
         serverless_config = None
-        instance_type = None
         if self.serverless:
             # For PyTorch or ChemProp we need at least 4GB of memory
             if needs_more_resources and mem_size < 4096:

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/model_script_utils/pytorch_utils.py

@@ -80,7 +80,6 @@ class TabularMLP(nn.Module):
         n_outputs: Number of output units
         task: "regression" or "classification"
         dropout: Dropout rate
-        use_batch_norm: Whether to use batch normalization
     """
 
     def __init__(
@@ -92,7 +91,6 @@ class TabularMLP(nn.Module):
         n_outputs: int,
         task: str = "regression",
         dropout: float = 0.1,
-        use_batch_norm: bool = True,
     ):
         super().__init__()
         self.task = task
@@ -112,8 +110,7 @@ class TabularMLP(nn.Module):
         layers = []
         for hidden_dim in hidden_layers:
             layers.append(nn.Linear(input_dim, hidden_dim))
-            if use_batch_norm:
-                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.LayerNorm(hidden_dim))
             layers.append(nn.LeakyReLU())
             layers.append(nn.Dropout(dropout))
             input_dim = hidden_dim
@@ -216,7 +213,6 @@ def create_model(
     n_outputs: int,
     task: str = "regression",
     dropout: float = 0.1,
-    use_batch_norm: bool = True,
 ) -> TabularMLP:
     """Create a TabularMLP model with appropriate embedding dimensions."""
     embedding_dims = compute_embedding_dims(categorical_cardinalities)
@@ -228,7 +224,6 @@ def create_model(
         n_outputs=n_outputs,
         task=task,
         dropout=dropout,
-        use_batch_norm=use_batch_norm,
     )
 
 
@@ -266,8 +261,8 @@ def train_model(
         train_dataset = TensorDataset(train_x_cont, dummy_cat, train_y)
         val_dataset = TensorDataset(val_x_cont, dummy_val_cat, val_y)
 
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
 
     # Loss and optimizer
     if task == "classification":
@@ -304,7 +299,7 @@ def train_model(
             out = model(x_cont, x_cat)
 
             if task == "classification":
-                loss = criterion(out, y.squeeze().long())
+                loss = criterion(out, y.view(-1).long())
             else:
                 loss = criterion(out, y)
 
@@ -323,7 +318,7 @@ def train_model(
                 out = model(x_cont, x_cat)
 
                 if task == "classification":
-                    loss = criterion(out, y.squeeze().long())
+                    loss = criterion(out, y.view(-1).long())
                 else:
                     loss = criterion(out, y)
                 val_losses.append(loss.item())
@@ -395,7 +390,6 @@ def load_model(path: str, device: str = "cpu") -> TabularMLP:
         n_outputs=config["n_outputs"],
         task=config["task"],
         dropout=config.get("dropout", 0.1),
-        use_batch_norm=config.get("use_batch_norm", True),
     )
 
     state_dict = torch.load(os.path.join(path, "model.pt"), map_location=device, weights_only=True)

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/model_scripts/chemprop/chemprop.template

@@ -34,7 +34,7 @@ DEFAULT_HYPERPARAMETERS = {
     "n_folds": 5,
     "max_epochs": 400,
     "patience": 50,
-    "batch_size": 32,
+    "batch_size": 64,
     # Message Passing (ignored when using foundation model)
     "hidden_dim": 700,
     "depth": 6,
@@ -590,8 +590,9 @@ if __name__ == "__main__":
             val_dataset.normalize_targets(target_scaler)
             output_transform = nn.UnscaleTransform.from_standard_scaler(target_scaler)
 
-        train_loader = data.build_dataloader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=3)
-        val_loader = data.build_dataloader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=3)
+        num_workers = min(os.cpu_count() or 4, 8)  # Scale with CPUs, cap at 8
+        train_loader = data.build_dataloader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, persistent_workers=num_workers > 0, pin_memory=True, prefetch_factor=2)
+        val_loader = data.build_dataloader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, persistent_workers=num_workers > 0, pin_memory=True, prefetch_factor=2)
 
         # Build model
         pl.seed_everything(hyperparameters["seed"] + fold_idx)
@@ -617,7 +618,7 @@ if __name__ == "__main__":
                 callbacks.append(pl.callbacks.ModelCheckpoint(
                     dirpath=args.model_dir, filename=f"best_{fold_idx}", monitor="val_loss", mode="min", save_top_k=1
                 ))
-            return pl.Trainer(accelerator="auto", max_epochs=max_epochs, logger=False, enable_progress_bar=True, callbacks=callbacks)
+            return pl.Trainer(accelerator="auto", max_epochs=max_epochs, precision="16-mixed", logger=False, enable_progress_bar=True, callbacks=callbacks)
 
         if use_two_phase:
             # Phase 1: Freeze MPNN, train FFN only
@@ -645,7 +646,7 @@ if __name__ == "__main__":
 
         # Out-of-fold predictions (using unscaled features - model's x_d_transform handles scaling)
         val_dps_raw, _ = _create_molecule_datapoints(df_val[smiles_column].tolist(), val_targets, val_extra_raw)
-        val_loader_pred = data.build_dataloader(data.MoleculeDataset(val_dps_raw), batch_size=batch_size, shuffle=False)
+        val_loader_pred = data.build_dataloader(data.MoleculeDataset(val_dps_raw), batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
 
         with torch.inference_mode():
             fold_preds = np.concatenate([p.numpy() for p in trainer.predict(mpnn, val_loader_pred)], axis=0)
@@ -706,7 +707,7 @@ if __name__ == "__main__":
         if len(ensemble_models) > 1:
             print("Computing prediction_std from ensemble...")
             val_dps, _ = _create_molecule_datapoints(df_val[smiles_column].tolist(), y_validate, val_extra_features)
-            val_loader = data.build_dataloader(data.MoleculeDataset(val_dps), batch_size=batch_size, shuffle=False)
+            val_loader = data.build_dataloader(data.MoleculeDataset(val_dps), batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
             trainer_pred = pl.Trainer(accelerator="auto", logger=False, enable_progress_bar=False)
 
             all_ens_preds = []

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/model_scripts/pytorch_model/pytorch.template

@@ -45,12 +45,11 @@ DEFAULT_HYPERPARAMETERS = {
     "n_folds": 5,
     "max_epochs": 200,
     "early_stopping_patience": 30,
-    "batch_size": 128,
+    "batch_size": 64,
     # Model architecture (larger capacity - ensemble provides regularization)
     "layers": "512-256-128",
     "learning_rate": 1e-3,
     "dropout": 0.05,
-    "use_batch_norm": True,
     # Loss function for regression (L1Loss=MAE, MSELoss=MSE, HuberLoss, SmoothL1Loss)
     "loss": "L1Loss",
     # Split strategy: "random", "scaffold", or "butina"
@@ -410,7 +409,6 @@ if __name__ == "__main__":
             n_outputs=n_outputs,
             task=task,
             dropout=hyperparameters["dropout"],
-            use_batch_norm=hyperparameters["use_batch_norm"],
         )
 
         # Train
@@ -517,7 +515,6 @@ if __name__ == "__main__":
         "n_outputs": n_outputs,
         "task": task,
         "dropout": hyperparameters["dropout"],
-        "use_batch_norm": hyperparameters["use_batch_norm"],
     }
 
     for idx, m in enumerate(ensemble_models):

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/model_scripts/pytorch_model/pytorch_utils.py

@@ -80,7 +80,6 @@ class TabularMLP(nn.Module):
         n_outputs: Number of output units
         task: "regression" or "classification"
         dropout: Dropout rate
-        use_batch_norm: Whether to use batch normalization
     """
 
     def __init__(
@@ -92,7 +91,6 @@ class TabularMLP(nn.Module):
         n_outputs: int,
         task: str = "regression",
         dropout: float = 0.1,
-        use_batch_norm: bool = True,
     ):
         super().__init__()
         self.task = task
@@ -112,8 +110,7 @@ class TabularMLP(nn.Module):
         layers = []
         for hidden_dim in hidden_layers:
             layers.append(nn.Linear(input_dim, hidden_dim))
-            if use_batch_norm:
-                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.LayerNorm(hidden_dim))
             layers.append(nn.LeakyReLU())
             layers.append(nn.Dropout(dropout))
             input_dim = hidden_dim
@@ -216,7 +213,6 @@ def create_model(
     n_outputs: int,
     task: str = "regression",
     dropout: float = 0.1,
-    use_batch_norm: bool = True,
 ) -> TabularMLP:
     """Create a TabularMLP model with appropriate embedding dimensions."""
     embedding_dims = compute_embedding_dims(categorical_cardinalities)
@@ -228,7 +224,6 @@ def create_model(
         n_outputs=n_outputs,
         task=task,
         dropout=dropout,
-        use_batch_norm=use_batch_norm,
     )
 
 
@@ -266,8 +261,8 @@ def train_model(
         train_dataset = TensorDataset(train_x_cont, dummy_cat, train_y)
         val_dataset = TensorDataset(val_x_cont, dummy_val_cat, val_y)
 
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
 
     # Loss and optimizer
     if task == "classification":
@@ -304,7 +299,7 @@ def train_model(
             out = model(x_cont, x_cat)
 
             if task == "classification":
-                loss = criterion(out, y.squeeze().long())
+                loss = criterion(out, y.view(-1).long())
             else:
                 loss = criterion(out, y)
 
@@ -323,7 +318,7 @@ def train_model(
                 out = model(x_cont, x_cat)
 
                 if task == "classification":
-                    loss = criterion(out, y.squeeze().long())
+                    loss = criterion(out, y.view(-1).long())
                 else:
                     loss = criterion(out, y)
                 val_losses.append(loss.item())
@@ -395,7 +390,6 @@ def load_model(path: str, device: str = "cpu") -> TabularMLP:
         n_outputs=config["n_outputs"],
         task=config["task"],
         dropout=config.get("dropout", 0.1),
-        use_batch_norm=config.get("use_batch_norm", True),
     )
 
     state_dict = torch.load(os.path.join(path, "model.pt"), map_location=device, weights_only=True)

{workbench-0.8.241 → workbench-0.8.246}/src/workbench/model_scripts/script_generation.py

@@ -3,6 +3,7 @@
 import os
 import shutil
 import logging
+import tempfile
 from pathlib import Path
 import importlib.util
 
@@ -41,7 +42,7 @@ def copy_imports_to_script_dir(script_path: str, imports: list[str]) -> None:
         print(f"Copied {source_path} to {destination_path}")
 
 
-def fill_template(template_path: str, params: dict, output_script: str) -> str:
+def fill_template(template_path: str, params: dict, output_script: str, output_dir: str = None) -> str:
     """
     Fill in the placeholders in the template with the values provided in params,
     ensuring that the correct Python data types are used.
@@ -49,6 +50,7 @@ def fill_template(template_path: str, params: dict, output_script: str) -> str:
         template_path (str): The path to the template file.
         params (dict): A dictionary with placeholder keys and their corresponding values.
         output_script (str): The name of the generated model script.
+        output_dir (str): The directory to write the output script to (default: same as template).
     Returns:
         str: The path to the generated model script.
     """
@@ -77,8 +79,10 @@ def fill_template(template_path: str, params: dict, output_script: str) -> str:
         log.critical(msg)
         raise ValueError(msg)
 
-    # Write out the generated model script and return the name
-    output_path = os.path.join(os.path.dirname(template_path), output_script)
+    # Write out the generated model script to output_dir (or template dir if not specified)
+    if output_dir is None:
+        output_dir = os.path.dirname(template_path)
+    output_path = os.path.join(output_dir, output_script)
     with open(output_path, "w") as fp:
         fp.write(template)
     return output_path
@@ -102,29 +106,29 @@ def generate_model_script(template_params: dict) -> str:
             - child_endpoints (list[str], optional): For META models, list of child endpoint names
 
     Returns:
-        str: The name of the generated model script
+        str: The path to the generated model script
     """
     from workbench.api import ModelType, ModelFramework  # Avoid circular import
 
     # Determine which template to use based on model type
     if template_params.get("model_class"):
         template_name = "scikit_learn.template"
-        model_script_dir = "scikit_learn"
+        model_script_dir_name = "scikit_learn"
     elif template_params["model_framework"] == ModelFramework.PYTORCH:
         template_name = "pytorch.template"
-        model_script_dir = "pytorch_model"
+        model_script_dir_name = "pytorch_model"
     elif template_params["model_framework"] == ModelFramework.CHEMPROP:
         template_name = "chemprop.template"
-        model_script_dir = "chemprop"
+        model_script_dir_name = "chemprop"
     elif template_params["model_framework"] == ModelFramework.META:
         template_name = "meta_model.template"
-        model_script_dir = "meta_model"
+        model_script_dir_name = "meta_model"
     elif template_params["model_type"] in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.CLASSIFIER]:
         template_name = "xgb_model.template"
-        model_script_dir = "xgb_model"
+        model_script_dir_name = "xgb_model"
     elif template_params["model_type"] == ModelType.ENSEMBLE_REGRESSOR:
         template_name = "ensemble_xgb.template"
-        model_script_dir = "ensemble_xgb"
+        model_script_dir_name = "ensemble_xgb"
     else:
         msg = f"ModelType: {template_params['model_type']} needs to set custom_script argument"
         log.critical(msg)
@@ -133,31 +137,42 @@ def generate_model_script(template_params: dict) -> str:
     # Model Type is an enumerated type, so we need to convert it to a string
     template_params["model_type"] = template_params["model_type"].value
 
-    # Load the template
-    dir_path = Path(__file__).parent.absolute()
-    model_script_dir = os.path.join(dir_path, model_script_dir)
-    template_path = os.path.join(model_script_dir, template_name)
+    # Load the template from the package directory
+    package_dir = Path(__file__).parent.absolute()
+    source_script_dir = os.path.join(package_dir, model_script_dir_name)
+    template_path = os.path.join(source_script_dir, template_name)
 
-    # Fill in the template and write out the generated model script
-    output_path = fill_template(template_path, template_params, "generated_model_script.py")
+    # Create a temp directory for the generated script and supporting files
+    # Note: This directory will be cleaned up by SageMaker after the training job
+    output_dir = tempfile.mkdtemp(prefix=f"workbench_{model_script_dir_name}_")
+    log.info(f"Generating model script in temp directory: {output_dir}")
 
-    # Report on any additional files in the model script directory
-    for file in os.listdir(model_script_dir):
-        if file not in ["generated_model_script.py", "requirements.txt"] and not file.endswith(".template"):
-            log.info(f"Additional file {file} found in model_script_dir...")
+    # Copy all supporting files (except templates and generated scripts) to the temp directory
+    for file in os.listdir(source_script_dir):
+        if file.endswith(".template") or file.startswith("generated_"):
+            continue
+        source_file = os.path.join(source_script_dir, file)
+        if os.path.isfile(source_file):
+            shutil.copy(source_file, output_dir)
+            log.info(f"Copied supporting file: {file}")
+
+    # Fill in the template and write the generated script to the temp directory
+    output_path = fill_template(template_path, template_params, "generated_model_script.py", output_dir)
 
     return output_path
 
 
 if __name__ == "__main__":
     """Exercise the Model Script Utilities"""
-    from workbench.api import ModelType
+    from workbench.api import ModelType, ModelFramework
 
     copy_imports_to_script_dir("/tmp/model.py", ["workbench.utils.chem_utils"])
 
     # Define the parameters for the model script (Classifier)
     my_params = {
         "model_type": ModelType.CLASSIFIER,
+        "model_framework": ModelFramework.XGBOOST,
+        "id_column": "id",
         "target_column": "wine_class",
         "feature_list": [
             "alcohol",
@@ -174,15 +189,18 @@ if __name__ == "__main__":
             "od280_od315_of_diluted_wines",
             "proline",
         ],
-        "model_metrics_s3_path": "s3://sandbox-workbench-artifacts/models/training/wine-classifier",
+        "model_metrics_s3_path": "s3://workbench-public-test-bucket/models/training/wine-classifier",
         "train_all_data": True,
+        "compressed_features": [],
+        "hyperparameters": {},
     }
     my_model_script = generate_model_script(my_params)
-    print(my_model_script)
+    print(f"Generated script: {my_model_script}")
 
     # Define the parameters for the model script (KMeans Clustering)
     my_params = {
         "model_type": ModelType.CLUSTERER,
+        "model_framework": ModelFramework.SKLEARN,
         "model_class": "KMeans",
         "model_imports": "from sklearn.cluster import KMeans",
         "target_column": None,
@@ -201,8 +219,8 @@ if __name__ == "__main__":
             "od280_od315_of_diluted_wines",
             "proline",
         ],
-        "model_metrics_s3_path": "s3://sandbox-workbench-artifacts/models/training/wine-clusters",
+        "model_metrics_s3_path": "s3://workbench-public-test-bucket/models/training/wine-clusters",
         "train_all_data": True,
     }
     my_model_script = generate_model_script(my_params)
-    print(my_model_script)
+    print(f"Generated script: {my_model_script}")

{workbench-0.8.241 → workbench-0.8.246}/src/workbench.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.241
+Version: 0.8.246
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License: MIT License

{workbench-0.8.241 → workbench-0.8.246}/src/workbench.egg-info/SOURCES.txt

@@ -576,7 +576,6 @@ src/workbench/model_script_utils/pytorch_utils.py
 src/workbench/model_script_utils/uq_harness.py
 src/workbench/model_scripts/script_generation.py
 src/workbench/model_scripts/chemprop/chemprop.template
-src/workbench/model_scripts/chemprop/generated_model_script.py
 src/workbench/model_scripts/chemprop/model_script_utils.py
 src/workbench/model_scripts/chemprop/requirements.txt
 src/workbench/model_scripts/custom_models/chem_info/Readme.md
@@ -602,19 +601,14 @@ src/workbench/model_scripts/custom_script_example/custom_model_script.py
 src/workbench/model_scripts/custom_script_example/requirements.txt
 src/workbench/model_scripts/ensemble_xgb/ensemble_xgb.template
 src/workbench/model_scripts/ensemble_xgb/requirements.txt
-src/workbench/model_scripts/meta_model/generated_model_script.py
 src/workbench/model_scripts/meta_model/meta_model.template
-src/workbench/model_scripts/pytorch_model/generated_model_script.py
 src/workbench/model_scripts/pytorch_model/model_script_utils.py
 src/workbench/model_scripts/pytorch_model/pytorch.template
 src/workbench/model_scripts/pytorch_model/pytorch_utils.py
 src/workbench/model_scripts/pytorch_model/requirements.txt
 src/workbench/model_scripts/pytorch_model/uq_harness.py
-src/workbench/model_scripts/scikit_learn/generated_model_script.py
 src/workbench/model_scripts/scikit_learn/requirements.txt
 src/workbench/model_scripts/scikit_learn/scikit_learn.template
-src/workbench/model_scripts/uq_models/generated_model_script.py
-src/workbench/model_scripts/xgb_model/generated_model_script.py
 src/workbench/model_scripts/xgb_model/model_script_utils.py
 src/workbench/model_scripts/xgb_model/requirements.txt
 src/workbench/model_scripts/xgb_model/uq_harness.py