workbench 0.8.180__tar.gz → 0.8.182__tar.gz

{workbench-0.8.180 → workbench-0.8.182}/PKG-INFO

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

{workbench-0.8.180 → workbench-0.8.182}/model_docker_images/ml_pipelines/Dockerfile

@@ -29,10 +29,10 @@ 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.177
+RUN pip install --no-cache-dir workbench==0.8.181
 
 # Add the script runner
 COPY ml_pipeline_runner.py /app/ml_pipeline_runner.py
 WORKDIR /app
 
-ENTRYPOINT ["python", "/app/ml_pipeline_runner.py"]
+ENTRYPOINT ["python", "/app/ml_pipeline_runner.py"]

{workbench-0.8.180 → workbench-0.8.182}/src/workbench/core/artifacts/endpoint_core.py

@@ -420,8 +420,12 @@ class EndpointCore(Artifact):
 
             # Capture the inference results and metrics
             if capture_name is not None:
+
+                # If we don't have an id_column, we'll pull it from the model's FeatureSet
+                if id_column is None:
+                    fs = FeatureSetCore(model.get_input())
+                    id_column = fs.id_column
                 description = capture_name.replace("_", " ").title()
-                features = model.features()
                 self._capture_inference_results(
                     capture_name, prediction_df, target_column, model_type, metrics, description, features, id_column
                 )
@@ -429,8 +433,6 @@ class EndpointCore(Artifact):
                 # For UQ Models we also capture the uncertainty metrics
                 if model_type in [ModelType.UQ_REGRESSOR]:
                     metrics = uq_metrics(prediction_df, target_column)
-
-                    # Now put into the Parameter Store Model Inference Namespace
                     self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
 
         # Return the prediction DataFrame
@@ -453,6 +455,55 @@ class EndpointCore(Artifact):
         cross_fold_metrics, out_of_fold_df = cross_fold_inference(model, nfolds=nfolds)
         if cross_fold_metrics:
             self.param_store.upsert(f"/workbench/models/{model.name}/inference/cross_fold", cross_fold_metrics)
+
+        # Capture the results
+        capture_name = "full_cross_fold"
+        description = capture_name.replace("_", " ").title()
+        target_column = model.target()
+        model_type = model.model_type
+
+        # Get the id_column from the model's FeatureSet
+        fs = FeatureSetCore(model.get_input())
+        id_column = fs.id_column
+
+        # Is this a UQ Model? If so, run full inference and merge the results
+        additional_columns = []
+        if model_type == ModelType.UQ_REGRESSOR:
+            self.log.important("UQ Regressor detected, running full inference to get uncertainty estimates...")
+
+            # Get the training view dataframe for inference
+            training_df = fs.view("training").pull_dataframe()
+
+            # Run inference on the endpoint to get UQ outputs
+            full_inference_df = self.inference(training_df)
+
+            # Identify UQ-specific columns (quantiles and prediction_std)
+            uq_columns = [col for col in full_inference_df.columns if col.startswith("q_") or col == "prediction_std"]
+
+            # Merge UQ columns with out-of-fold predictions
+            if uq_columns:
+                # Keep id_column and UQ columns, drop 'prediction' to avoid conflict
+                merge_columns = [id_column] + uq_columns
+                uq_df = full_inference_df[merge_columns]
+
+                out_of_fold_df = pd.merge(out_of_fold_df, uq_df, on=id_column, how="left")
+                additional_columns = uq_columns
+                self.log.info(f"Added UQ columns: {', '.join(additional_columns)}")
+
+                # Also compute UQ metrics
+                metrics = uq_metrics(out_of_fold_df, target_column)
+                self.param_store.upsert(f"/workbench/models/{model.name}/inference/{capture_name}", metrics)
+
+        self._capture_inference_results(
+            capture_name,
+            out_of_fold_df,
+            target_column,
+            model_type,
+            pd.DataFrame([cross_fold_metrics["summary_metrics"]]),
+            description,
+            features=additional_columns,
+            id_column=id_column,
+        )
         return cross_fold_metrics, out_of_fold_df
 
     def fast_inference(self, eval_df: pd.DataFrame, threads: int = 4) -> pd.DataFrame:
@@ -648,6 +699,10 @@ class EndpointCore(Artifact):
     @staticmethod
     def _hash_dataframe(df: pd.DataFrame, hash_length: int = 8):
         # Internal: Compute a data hash for the dataframe
+        if df.empty:
+            return "--hash--"
+
+        # Sort the dataframe by columns to ensure consistent ordering
         df = df.copy()
         df = df.sort_values(by=sorted(df.columns.tolist()))
         row_hashes = pd.util.hash_pandas_object(df, index=False)
@@ -702,8 +757,8 @@ class EndpointCore(Artifact):
         wr.s3.to_csv(metrics, f"{inference_capture_path}/inference_metrics.csv", index=False)
 
         # Grab the target column, prediction column, any _proba columns, and the ID column (if present)
-        prediction_col = "prediction" if "prediction" in pred_results_df.columns else "predictions"
-        output_columns = [target_column, prediction_col]
+        output_columns = [target_column]
+        output_columns += [col for col in pred_results_df.columns if "prediction" in col]
 
         # Add any _proba columns to the output columns
         output_columns += [col for col in pred_results_df.columns if col.endswith("_proba")]
@@ -713,7 +768,7 @@ class EndpointCore(Artifact):
 
         # Add the ID column
         if id_column and id_column in pred_results_df.columns:
-            output_columns.append(id_column)
+            output_columns.insert(0, id_column)
 
         # Write the predictions to our S3 Model Inference Folder
         self.log.info(f"Writing predictions to {inference_capture_path}/inference_predictions.csv")
@@ -1134,6 +1189,10 @@ if __name__ == "__main__":
     # Run predictions using the fast_inference method
     fast_results = my_endpoint.fast_inference(my_eval_df)
 
+    # Test the cross_fold_inference method
+    print("Running Cross-Fold Inference...")
+    metrics, all_results = my_endpoint.cross_fold_inference()
+
     # Run Inference and metrics for a Classification Endpoint
     class_endpoint = EndpointCore("wine-classification")
     auto_predictions = class_endpoint.auto_inference()
@@ -1142,6 +1201,10 @@ if __name__ == "__main__":
     target = "wine_class"
     print(class_endpoint.generate_confusion_matrix(target, auto_predictions))
 
+    # Test the cross_fold_inference method
+    print("Running Cross-Fold Inference...")
+    metrics, all_results = class_endpoint.cross_fold_inference()
+
     # Test the class method delete (commented out for now)
     # from workbench.api import Model
     # model = Model("abalone-regression")

{workbench-0.8.180 → workbench-0.8.182}/src/workbench/model_scripts/custom_models/uq_models/generated_model_script.py

@@ -18,11 +18,11 @@ from typing import List, Tuple
 
 # Template Placeholders
 TEMPLATE_PARAMS = {
-    "target": "solubility",
-    "features": ['molwt', 'mollogp', 'molmr', 'heavyatomcount', 'numhacceptors', 'numhdonors', 'numheteroatoms', 'numrotatablebonds', 'numvalenceelectrons', 'numaromaticrings', 'numsaturatedrings', 'numaliphaticrings', 'ringcount', 'tpsa', 'labuteasa', 'balabanj', 'bertzct'],
+    "target": "udm_asy_res_free_percent",
+    "features": ['vsa_estate6', 'naromatom', 'mollogp', 'fr_nh2', 'mp', 'c2sp2', 'xch_3d', 'axp_6d', 'bcut2d_mrhi', 'fr_benzene', 'mz', 'slogp_vsa6', 'fr_halogen', 'bcut2d_mwhi', 'vsa_estate4', 'slogp_vsa3', 'estate_vsa5', 'minestateindex', 'axp_3dv', 'estate_vsa3', 'vsa_estate9', 'molwt', 'hallkieralpha', 'fpdensitymorgan1', 'peoe_vsa13', 'xpc_5d', 'bcut2d_chghi', 'peoe_vsa8', 'axp_0dv', 'axp_2d', 'chi2v', 'bcut2d_logphi', 'axp_5d', 'peoe_vsa2', 'estate_vsa6', 'qed', 'numrotatablebonds', 'xc_3dv', 'peoe_vsa3', 'balabanj', 'slogp_vsa5', 'mv', 'vsa_estate2', 'bcut2d_mwlow', 'xch_7d', 'chi3n', 'vsa_estate8', 'estate_vsa4', 'xp_7dv', 'fr_nh1', 'vsa_estate3', 'fr_ketone_topliss', 'minpartialcharge', 'phi', 'peoe_vsa10', 'vsa_estate7', 'estate_vsa7', 'tpsa', 'kappa3', 'kappa2', 'bcut2d_logplow', 'xch_6d', 'maxpartialcharge', 'vsa_estate1', 'peoe_vsa9', 'axp_1d', 'fr_ar_n', 'chi2n', 'vsa_estate5', 'xp_4dv', 'slogp_vsa10', 'num_stereobonds', 'peoe_vsa11', 'bcut2d_chglo', 'chi1v', 'peoe_vsa7', 'bertzct', 'axp_2dv', 'estate_vsa2', 'smr_vsa9', 'peoe_vsa6', 'num_s_centers', 'num_r_centers', 'xch_7dv', 'xc_5d', 'axp_4dv', 'xc_5dv', 'mi', 'xc_3d', 'fpdensitymorgan2', 'xp_0dv', 'nhohcount', 'numatomstereocenters', 'mse', 'smr_vsa3', 'peoe_vsa12', 'nocount', 'fpdensitymorgan3', 'minabsestateindex', 'bcut2d_mrlow', 'axp_5dv', 'sz', 'vsa_estate10', 'axp_3d', 'xch_6dv', 'xch_4d', 'xc_6d', 'estate_vsa8', 'mpe', 'smr_vsa7', 'numhdonors', 'smr_vsa1', 'xp_5d', 'fr_para_hydroxylation', 'chi3v', 'xpc_6dv', 'nbase', 'heavyatommolwt', 'avgipc', 'maxestateindex', 'smr_vsa6', 'fr_bicyclic', 'xc_4dv', 'xp_7d', 'smr_vsa5', 'xpc_4d', 'smr_vsa4', 'peoe_vsa4', 'numheteroatoms', 'fr_nhpyrrole', 'axp_4d', 'smr_vsa10', 'xp_6d', 'sps', 'mare', 'slogp_vsa2', 'axp_0d', 'slogp_vsa4', 'fr_al_oh', 'numheterocycles', 'labuteasa', 'xp_3d', 'chi4n', 'fractioncsp3', 'maxabspartialcharge', 'fr_al_oh_notert', 'peoe_vsa1', 'axp_7dv', 'slogp_vsa11', 'peoe_vsa5', 'xpc_5dv', 'xpc_6d', 'xp_2d', 'xp_3dv', 'fr_ndealkylation1', 'axp_7d', 'estate_vsa9', 'molmr', 'num_stereocenters', 'si', 'estate_vsa1', 'xc_6dv', 'chi0v', 'fr_oxazole', 'axp_6dv', 'xp_6dv', 'xp_4d', 'numaliphaticheterocycles', 'fr_imine', 'fr_imidazole', 'xp_5dv', 'fr_piperdine', 'slogp_vsa7', 'chi1', 'c1sp2', 'numaromaticheterocycles', 'xpc_4dv', 'c3sp2', 'fr_aniline', 'fr_piperzine', 'axp_1dv', 'xch_4dv', 'chi4v', 'chi1n', 'minabspartialcharge', 'slogp_vsa1', 'fr_nh0', 'chi0n', 'c2sp3', 'xc_4d', 'xch_5dv', 'peoe_vsa14', 'xch_5d', 'numsaturatedrings', 'fr_pyridine', 'kappa1', 'slogp_vsa8', 'xp_2dv', 'fr_ar_coo', 'numvalenceelectrons'],
     "compressed_features": [],
-    "train_all_data": False,
-    "hyperparameters": {'n_estimators': 200, 'max_depth': 6, 'learning_rate': 0.05, 'subsample': 0.7, 'colsample_bytree': 0.3, 'colsample_bylevel': 0.5, 'min_child_weight': 5, 'gamma': 0.2, 'reg_alpha': 0.5, 'reg_lambda': 2.0, 'scale_pos_weight': 1},
+    "train_all_data": True,
+    "hyperparameters": {},
 }
 
 
@@ -303,7 +303,7 @@ if __name__ == "__main__":
         print(f"  {conf_level * 100:.0f}% CI: Mean width={np.mean(widths):.3f}, Std={np.std(widths):.3f}")
 
     # Save the trained XGBoost model
-    xgb_model.save_model(os.path.join(args.model_dir, "xgb_model.json"))
+    joblib.dump(xgb_model, os.path.join(args.model_dir, "xgb_model.joblib"))
 
     # Save all MAPIE models
     for model_name, model in mapie_models.items():
@@ -349,9 +349,8 @@ def model_fn(model_dir) -> dict:
         config = json.load(fp)
 
     # Load XGBoost regressor
-    xgb_path = os.path.join(model_dir, "xgb_model.json")
-    xgb_model = XGBRegressor(enable_categorical=True)
-    xgb_model.load_model(xgb_path)
+    xgb_path = os.path.join(model_dir, "xgb_model.joblib")
+    xgb_model = joblib.load(xgb_path)
 
     # Load all MAPIE models
     mapie_models = {}

{workbench-0.8.180 → workbench-0.8.182}/src/workbench/model_scripts/script_generation.py

@@ -117,8 +117,8 @@ def generate_model_script(template_params: dict) -> str:
         template_name = "xgb_model.template"
         model_script_dir = "xgb_model"
     elif template_params["model_type"] == ModelType.UQ_REGRESSOR:
-        template_name = "quant_regression.template"
-        model_script_dir = "quant_regression"
+        template_name = "mapie.template"
+        model_script_dir = "uq_models"
     elif template_params["model_type"] == ModelType.ENSEMBLE_REGRESSOR:
         template_name = "ensemble_xgb.template"
         model_script_dir = "ensemble_xgb"