workbench 0.8.213__py3-none-any.whl → 0.8.219__py3-none-any.whl

workbench/model_scripts/meta_model/generated_model_script.py

@@ -0,0 +1,209 @@
+# Meta Model Template for Workbench
+#
+# NOTE: This is called a "meta model" but it's really a "meta endpoint" - it aggregates
+# predictions from multiple child endpoints. We call it a "model" because Workbench
+# creates Model artifacts that get deployed as Endpoints, so this follows that pattern.
+#
+# Assumptions/Shortcuts:
+# - All child endpoints are regression models
+# - All child endpoints output 'prediction' and 'confidence' columns
+# - Aggregation uses model weights (provided at meta model creation time)
+#
+# This template:
+# - Has no real training phase (just saves metadata including model weights)
+# - At inference time, calls child endpoints and aggregates their predictions
+
+import argparse
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import StringIO
+
+import pandas as pd
+
+from workbench_bridges.endpoints.fast_inference import fast_inference
+
+# Template parameters (filled in by Workbench)
+TEMPLATE_PARAMS = {
+    "child_endpoints": ['logd-reg-pytorch', 'logd-reg-chemprop'],
+    "target_column": "logd",
+    "model_weights": {'logd-reg-pytorch': 0.4228205813233993, 'logd-reg-chemprop': 0.5771794186766008},
+    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/logd-meta/training",
+    "aws_region": "us-west-2",
+}
+
+
+def invoke_endpoints_parallel(endpoint_names: list[str], df: pd.DataFrame) -> dict[str, pd.DataFrame]:
+    """Call multiple child endpoints in parallel and collect their results.
+
+    Args:
+        endpoint_names: List of endpoint names to call
+        df: Input DataFrame to send to each endpoint
+
+    Returns:
+        Dict mapping endpoint_name -> result DataFrame (or None if failed)
+    """
+    results = {}
+
+    def call_endpoint(name: str) -> tuple[str, pd.DataFrame | None]:
+        try:
+            return name, fast_inference(name, df)
+        except Exception as e:
+            print(f"Error calling endpoint {name}: {e}")
+            return name, None
+
+    with ThreadPoolExecutor(max_workers=len(endpoint_names)) as executor:
+        futures = {executor.submit(call_endpoint, name): name for name in endpoint_names}
+        for future in as_completed(futures):
+            name, result = future.result()
+            results[name] = result
+
+    return results
+
+
+def aggregate_predictions(results: dict[str, pd.DataFrame], model_weights: dict[str, float]) -> pd.DataFrame:
+    """Aggregate predictions from multiple endpoints using model weights.
+
+    Args:
+        results: Dict mapping endpoint_name -> predictions DataFrame
+                 Each DataFrame must have 'prediction' and 'confidence' columns
+        model_weights: Dict mapping endpoint_name -> weight
+
+    Returns:
+        DataFrame with aggregated prediction, prediction_std, and confidence
+    """
+    # Filter out failed endpoints
+    valid_results = {k: v for k, v in results.items() if v is not None}
+    if not valid_results:
+        raise ValueError("All child endpoints failed")
+
+    # Use first result as base (for id columns, etc.)
+    first_df = list(valid_results.values())[0]
+    output_df = first_df.drop(columns=["prediction", "confidence", "prediction_std"], errors="ignore").copy()
+
+    # Build DataFrames of predictions and confidences from all endpoints
+    pred_df = pd.DataFrame({name: df["prediction"] for name, df in valid_results.items()})
+    conf_df = pd.DataFrame({name: df["confidence"] for name, df in valid_results.items()})
+
+    # Apply model weights (renormalize for valid endpoints only)
+    valid_weights = {k: model_weights.get(k, 1.0) for k in valid_results}
+    weight_sum = sum(valid_weights.values())
+    normalized_weights = {k: v / weight_sum for k, v in valid_weights.items()}
+
+    # Weighted average
+    output_df["prediction"] = sum(pred_df[name] * w for name, w in normalized_weights.items())
+
+    # Ensemble std across child endpoints
+    output_df["prediction_std"] = pred_df.std(axis=1)
+
+    # Aggregated confidence: weighted mean of child confidences
+    output_df["confidence"] = sum(conf_df[name] * w for name, w in normalized_weights.items())
+
+    return output_df
+
+
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load meta model configuration."""
+    with open(os.path.join(model_dir, "meta_config.json")) as f:
+        config = json.load(f)
+
+    # Set AWS_REGION for fast_inference (baked in at training time)
+    if config.get("aws_region"):
+        os.environ["AWS_REGION"] = config["aws_region"]
+
+    print(f"Meta model loaded: {len(config['child_endpoints'])} child endpoints")
+    print(f"Model weights: {config.get('model_weights')}")
+    print(f"AWS region: {config.get('aws_region')}")
+    return config
+
+
+def input_fn(input_data, content_type):
+    """Parse input data and return a DataFrame."""
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    # Decode bytes to string if necessary
+    if isinstance(input_data, bytes):
+        input_data = input_data.decode("utf-8")
+
+    if "text/csv" in content_type:
+        return pd.read_csv(StringIO(input_data))
+    elif "application/json" in content_type:
+        return pd.DataFrame(json.loads(input_data))
+    else:
+        raise ValueError(f"{content_type} not supported!")
+
+
+def output_fn(output_df, accept_type):
+    """Supports both CSV and JSON output formats."""
+    if "text/csv" in accept_type:
+        return output_df.to_csv(index=False), "text/csv"
+    elif "application/json" in accept_type:
+        return output_df.to_json(orient="records"), "application/json"
+    else:
+        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, config: dict) -> pd.DataFrame:
+    """Run inference by calling child endpoints and aggregating results."""
+    child_endpoints = config["child_endpoints"]
+    model_weights = config.get("model_weights", {})
+
+    print(f"Calling {len(child_endpoints)} child endpoints: {child_endpoints}")
+
+    # Call all child endpoints
+    results = invoke_endpoints_parallel(child_endpoints, df)
+
+    # Report status
+    for name, result in results.items():
+        status = f"{len(result)} rows" if result is not None else "FAILED"
+        print(f"  {name}: {status}")
+
+    # Aggregate predictions using model weights
+    output_df = aggregate_predictions(results, model_weights)
+
+    print(f"Aggregated {len(output_df)} predictions from {len(results)} endpoints")
+    return output_df
+
+
+# =============================================================================
+# Training (just saves configuration - no actual training)
+# =============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    parser.add_argument("--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    args = parser.parse_args()
+
+    child_endpoints = TEMPLATE_PARAMS["child_endpoints"]
+    target_column = TEMPLATE_PARAMS["target_column"]
+    model_weights = TEMPLATE_PARAMS["model_weights"]
+    aws_region = TEMPLATE_PARAMS["aws_region"]
+
+    print("=" * 60)
+    print("Meta Model Configuration")
+    print("=" * 60)
+    print(f"Child endpoints: {child_endpoints}")
+    print(f"Target column: {target_column}")
+    print(f"Model weights: {model_weights}")
+    print(f"AWS region: {aws_region}")
+
+    # Save configuration for inference
+    config = {
+        "child_endpoints": child_endpoints,
+        "target_column": target_column,
+        "model_weights": model_weights,
+        "aws_region": aws_region,
+    }
+
+    with open(os.path.join(args.model_dir, "meta_config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    print(f"\nMeta model configuration saved to {args.model_dir}")

workbench/model_scripts/meta_model/meta_model.template

@@ -0,0 +1,209 @@
+# Meta Model Template for Workbench
+#
+# NOTE: This is called a "meta model" but it's really a "meta endpoint" - it aggregates
+# predictions from multiple child endpoints. We call it a "model" because Workbench
+# creates Model artifacts that get deployed as Endpoints, so this follows that pattern.
+#
+# Assumptions/Shortcuts:
+# - All child endpoints are regression models
+# - All child endpoints output 'prediction' and 'confidence' columns
+# - Aggregation uses model weights (provided at meta model creation time)
+#
+# This template:
+# - Has no real training phase (just saves metadata including model weights)
+# - At inference time, calls child endpoints and aggregates their predictions
+
+import argparse
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import StringIO
+
+import pandas as pd
+
+from workbench_bridges.endpoints.fast_inference import fast_inference
+
+# Template parameters (filled in by Workbench)
+TEMPLATE_PARAMS = {
+    "child_endpoints": "{{child_endpoints}}",
+    "target_column": "{{target_column}}",
+    "model_weights": "{{model_weights}}",
+    "model_metrics_s3_path": "{{model_metrics_s3_path}}",
+    "aws_region": "{{aws_region}}",
+}
+
+
+def invoke_endpoints_parallel(endpoint_names: list[str], df: pd.DataFrame) -> dict[str, pd.DataFrame]:
+    """Call multiple child endpoints in parallel and collect their results.
+
+    Args:
+        endpoint_names: List of endpoint names to call
+        df: Input DataFrame to send to each endpoint
+
+    Returns:
+        Dict mapping endpoint_name -> result DataFrame (or None if failed)
+    """
+    results = {}
+
+    def call_endpoint(name: str) -> tuple[str, pd.DataFrame | None]:
+        try:
+            return name, fast_inference(name, df)
+        except Exception as e:
+            print(f"Error calling endpoint {name}: {e}")
+            return name, None
+
+    with ThreadPoolExecutor(max_workers=len(endpoint_names)) as executor:
+        futures = {executor.submit(call_endpoint, name): name for name in endpoint_names}
+        for future in as_completed(futures):
+            name, result = future.result()
+            results[name] = result
+
+    return results
+
+
+def aggregate_predictions(results: dict[str, pd.DataFrame], model_weights: dict[str, float]) -> pd.DataFrame:
+    """Aggregate predictions from multiple endpoints using model weights.
+
+    Args:
+        results: Dict mapping endpoint_name -> predictions DataFrame
+                 Each DataFrame must have 'prediction' and 'confidence' columns
+        model_weights: Dict mapping endpoint_name -> weight
+
+    Returns:
+        DataFrame with aggregated prediction, prediction_std, and confidence
+    """
+    # Filter out failed endpoints
+    valid_results = {k: v for k, v in results.items() if v is not None}
+    if not valid_results:
+        raise ValueError("All child endpoints failed")
+
+    # Use first result as base (for id columns, etc.)
+    first_df = list(valid_results.values())[0]
+    output_df = first_df.drop(columns=["prediction", "confidence", "prediction_std"], errors="ignore").copy()
+
+    # Build DataFrames of predictions and confidences from all endpoints
+    pred_df = pd.DataFrame({name: df["prediction"] for name, df in valid_results.items()})
+    conf_df = pd.DataFrame({name: df["confidence"] for name, df in valid_results.items()})
+
+    # Apply model weights (renormalize for valid endpoints only)
+    valid_weights = {k: model_weights.get(k, 1.0) for k in valid_results}
+    weight_sum = sum(valid_weights.values())
+    normalized_weights = {k: v / weight_sum for k, v in valid_weights.items()}
+
+    # Weighted average
+    output_df["prediction"] = sum(pred_df[name] * w for name, w in normalized_weights.items())
+
+    # Ensemble std across child endpoints
+    output_df["prediction_std"] = pred_df.std(axis=1)
+
+    # Aggregated confidence: weighted mean of child confidences
+    output_df["confidence"] = sum(conf_df[name] * w for name, w in normalized_weights.items())
+
+    return output_df
+
+
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load meta model configuration."""
+    with open(os.path.join(model_dir, "meta_config.json")) as f:
+        config = json.load(f)
+
+    # Set AWS_REGION for fast_inference (baked in at training time)
+    if config.get("aws_region"):
+        os.environ["AWS_REGION"] = config["aws_region"]
+
+    print(f"Meta model loaded: {len(config['child_endpoints'])} child endpoints")
+    print(f"Model weights: {config.get('model_weights')}")
+    print(f"AWS region: {config.get('aws_region')}")
+    return config
+
+
+def input_fn(input_data, content_type):
+    """Parse input data and return a DataFrame."""
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    # Decode bytes to string if necessary
+    if isinstance(input_data, bytes):
+        input_data = input_data.decode("utf-8")
+
+    if "text/csv" in content_type:
+        return pd.read_csv(StringIO(input_data))
+    elif "application/json" in content_type:
+        return pd.DataFrame(json.loads(input_data))
+    else:
+        raise ValueError(f"{content_type} not supported!")
+
+
+def output_fn(output_df, accept_type):
+    """Supports both CSV and JSON output formats."""
+    if "text/csv" in accept_type:
+        return output_df.to_csv(index=False), "text/csv"
+    elif "application/json" in accept_type:
+        return output_df.to_json(orient="records"), "application/json"
+    else:
+        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, config: dict) -> pd.DataFrame:
+    """Run inference by calling child endpoints and aggregating results."""
+    child_endpoints = config["child_endpoints"]
+    model_weights = config.get("model_weights", {})
+
+    print(f"Calling {len(child_endpoints)} child endpoints: {child_endpoints}")
+
+    # Call all child endpoints
+    results = invoke_endpoints_parallel(child_endpoints, df)
+
+    # Report status
+    for name, result in results.items():
+        status = f"{len(result)} rows" if result is not None else "FAILED"
+        print(f"  {name}: {status}")
+
+    # Aggregate predictions using model weights
+    output_df = aggregate_predictions(results, model_weights)
+
+    print(f"Aggregated {len(output_df)} predictions from {len(results)} endpoints")
+    return output_df
+
+
+# =============================================================================
+# Training (just saves configuration - no actual training)
+# =============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    parser.add_argument("--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    args = parser.parse_args()
+
+    child_endpoints = TEMPLATE_PARAMS["child_endpoints"]
+    target_column = TEMPLATE_PARAMS["target_column"]
+    model_weights = TEMPLATE_PARAMS["model_weights"]
+    aws_region = TEMPLATE_PARAMS["aws_region"]
+
+    print("=" * 60)
+    print("Meta Model Configuration")
+    print("=" * 60)
+    print(f"Child endpoints: {child_endpoints}")
+    print(f"Target column: {target_column}")
+    print(f"Model weights: {model_weights}")
+    print(f"AWS region: {aws_region}")
+
+    # Save configuration for inference
+    config = {
+        "child_endpoints": child_endpoints,
+        "target_column": target_column,
+        "model_weights": model_weights,
+        "aws_region": aws_region,
+    }
+
+    with open(os.path.join(args.model_dir, "meta_config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    print(f"\nMeta model configuration saved to {args.model_dir}")

workbench/model_scripts/pytorch_model/generated_model_script.py

@@ -5,51 +5,36 @@
 # - Out-of-fold predictions for validation metrics
 # - Categorical feature embedding via TabularMLP
 # - Compressed feature decompression
+#
+# NOTE: Imports are structured to minimize serverless endpoint startup time.
+# Heavy imports (sklearn, awswrangler) are deferred to training time.
 
-import argparse
 import json
 import os
 
-import awswrangler as wr
 import joblib
 import numpy as np
 import pandas as pd
 import torch
-from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
-from sklearn.preprocessing import LabelEncoder
-
-# Enable Tensor Core optimization for GPUs that support it
-torch.set_float32_matmul_precision("medium")
 
 from model_script_utils import (
-    check_dataframe,
-    compute_classification_metrics,
-    compute_regression_metrics,
     convert_categorical_types,
     decompress_features,
     expand_proba_column,
     input_fn,
     match_features_case_insensitive,
     output_fn,
-    print_classification_metrics,
-    print_confusion_matrix,
-    print_regression_metrics,
 )
 from pytorch_utils import (
     FeatureScaler,
-    create_model,
     load_model,
     predict,
     prepare_data,
-    save_model,
-    train_model,
 )
 from uq_harness import (
     compute_confidence,
     load_uq_models,
     predict_intervals,
-    save_uq_models,
-    train_uq_models,
 )
 
 # =============================================================================
@@ -59,13 +44,15 @@ DEFAULT_HYPERPARAMETERS = {
     # Training parameters
     "n_folds": 5,
     "max_epochs": 200,
-    "early_stopping_patience": 20,
+    "early_stopping_patience": 30,
     "batch_size": 128,
-    # Model architecture
-    "layers": "256-128-64",
+    # Model architecture (larger capacity - ensemble provides regularization)
+    "layers": "512-256-128",
     "learning_rate": 1e-3,
-    "dropout": 0.1,
+    "dropout": 0.05,
     "use_batch_norm": True,
+    # Loss function for regression (L1Loss=MAE, MSELoss=MSE, HuberLoss, SmoothL1Loss)
+    "loss": "L1Loss",
     # Random seed
     "seed": 42,
 }
@@ -74,10 +61,10 @@ DEFAULT_HYPERPARAMETERS = {
 TEMPLATE_PARAMS = {
     "model_type": "uq_regressor",
     "target": "udm_asy_res_efflux_ratio",
-    "features": ['smr_vsa4', 'tpsa', 'numhdonors', 'nhohcount', 'nbase', 'vsa_estate3', 'fr_guanido', 'mollogp', 'peoe_vsa8', 'peoe_vsa1', 'fr_imine', 'vsa_estate2', 'estate_vsa10', 'asphericity', 'xc_3dv', 'smr_vsa3', 'charge_centroid_distance', 'c3sp3', 'nitrogen_span', 'estate_vsa2', 'minpartialcharge', 'hba_hbd_ratio', 'slogp_vsa1', 'axp_7d', 'nocount', 'vsa_estate4', 'vsa_estate6', 'estate_vsa4', 'xc_4dv', 'xc_4d', 'num_s_centers', 'vsa_estate9', 'chi2v', 'axp_5d', 'mi', 'mse', 'bcut2d_mrhi', 'smr_vsa6', 'hallkieralpha', 'balabanj', 'amphiphilic_moment', 'type_ii_pattern_count', 'minabsestateindex', 'bcut2d_mwlow', 'axp_0dv', 'slogp_vsa5', 'axp_2d', 'axp_1dv', 'xch_5d', 'peoe_vsa10', 'molecular_asymmetry', 'kappa3', 'estate_vsa3', 'sse', 'bcut2d_logphi', 'fr_imidazole', 'molecular_volume_3d', 'bertzct', 'maxestateindex', 'aromatic_interaction_score', 'axp_3d', 'radius_of_gyration', 'vsa_estate7', 'si', 'axp_5dv', 'molecular_axis_length', 'estate_vsa6', 'fpdensitymorgan1', 'axp_6d', 'estate_vsa9', 'fpdensitymorgan2', 'xp_0dv', 'xp_6dv', 'molmr', 'qed', 'estate_vsa8', 'peoe_vsa9', 'xch_6dv', 'xp_7d', 'slogp_vsa2', 'xp_5dv', 'bcut2d_chghi', 'xch_6d', 'chi0n', 'slogp_vsa3', 'chi1v', 'chi3v', 'bcut2d_chglo', 'axp_1d', 'mp', 'num_defined_stereocenters', 'xp_3dv', 'bcut2d_mrlow', 'fr_al_oh', 'peoe_vsa7', 'chi2n', 'axp_6dv', 'axp_2dv', 'chi4n', 'xc_3d', 'axp_7dv', 'vsa_estate8', 'xch_7d', 'maxpartialcharge', 'chi1n', 'peoe_vsa2', 'axp_3dv', 'bcut2d_logplow', 'mv', 'xpc_5dv', 'kappa2', 'vsa_estate5', 'xp_5d', 'mm', 'maxabspartialcharge', 'axp_4dv', 'maxabsestateindex', 'axp_4d', 'xch_4dv', 'xp_2dv', 'heavyatommolwt', 'numatomstereocenters', 'xp_7dv', 'numsaturatedheterocycles', 'xp_3d', 'kappa1', 'mz', 'axp_0d', 'chi1', 'xch_4d', 'smr_vsa1', 'xp_2d', 'estate_vsa5', 'phi', 'fr_ether', 'xc_5d', 'c1sp3', 'estate_vsa7', 'estate_vsa1', 'vsa_estate1', 'slogp_vsa4', 'avgipc', 'smr_vsa10', 'numvalenceelectrons', 'xc_5dv', 'peoe_vsa12', 'peoe_vsa6', 'xpc_5d', 'xpc_6d', 'minestateindex', 'chi3n', 'smr_vsa5', 'xp_4d', 'numheteroatoms', 'fpdensitymorgan3', 'xpc_4d', 'sps', 'xp_1d', 'sv', 'fr_ar_n', 'slogp_vsa10', 'c2sp3', 'xpc_4dv', 'chi0v', 'xpc_6dv', 'xp_1dv', 'vsa_estate10', 'sare', 'c2sp2', 'mpe', 'xch_7dv', 'chi4v', 'type_i_pattern_count', 'sp', 'slogp_vsa8', 'amide_count', 'num_stereocenters', 'num_r_centers', 'tertiary_amine_count', 'spe', 'xp_4dv', 'numsaturatedrings', 'mare', 'numhacceptors', 'chi0', 'fractioncsp3', 'fr_nh0', 'xch_5dv', 'fr_aniline', 'smr_vsa7', 'labuteasa', 'c3sp2', 'xp_0d', 'xp_6d', 'peoe_vsa11', 'fr_ar_nh', 'molwt', 'intramolecular_hbond_potential', 'peoe_vsa3', 'fr_nhpyrrole', 'numaliphaticrings', 'hybratio', 'smr_vsa9', 'peoe_vsa13', 'bcut2d_mwhi', 'c1sp2', 'slogp_vsa11', 'numrotatablebonds', 'numaliphaticcarbocycles', 'slogp_vsa6', 'peoe_vsa4', 'numunspecifiedatomstereocenters', 'xc_6d', 'xc_6dv', 'num_unspecified_stereocenters', 'sz', 'minabspartialcharge', 'fcsp3', 'c1sp1', 'fr_piperzine', 'numaliphaticheterocycles', 'numamidebonds', 'fr_benzene', 'numaromaticheterocycles', 'sm', 'fr_priamide', 'fr_piperdine', 'fr_methoxy', 'c4sp3', 'fr_c_o_nocoo', 'exactmolwt', 'stereo_complexity', 'fr_hoccn', 'numaromaticcarbocycles', 'fr_nh2', 'numheterocycles', 'fr_morpholine', 'fr_ketone', 'fr_nh1', 'frac_defined_stereo', 'fr_aryl_methyl', 'fr_alkyl_halide', 'fr_phenol', 'fr_al_oh_notert', 'fr_ar_oh', 'fr_pyridine', 'fr_amide', 'slogp_vsa7', 'fr_halogen', 'numsaturatedcarbocycles', 'slogp_vsa12', 'fr_ndealkylation1', 'xch_3d', 'fr_bicyclic', 'naromatom', 'narombond'],
+    "features": ['fingerprint'],
     "id_column": "udm_mol_bat_id",
-    "compressed_features": [],
-    "model_metrics_s3_path": "s3://ideaya-sageworks-bucket/models/caco2-er-reg-pytorch/training",
+    "compressed_features": ['fingerprint'],
+    "model_metrics_s3_path": "s3://ideaya-sageworks-bucket/models/caco2-er-reg-fp-pytorch/training",
     "hyperparameters": {},
 }
 
@@ -86,7 +73,7 @@ TEMPLATE_PARAMS = {
 # Model Loading (for SageMaker inference)
 # =============================================================================
 def model_fn(model_dir: str) -> dict:
-    """Load TabularMLP ensemble from the specified directory."""
+    """Load PyTorch TabularMLP ensemble from the specified directory."""
     # Load ensemble metadata
     metadata_path = os.path.join(model_dir, "ensemble_metadata.joblib")
     if os.path.exists(metadata_path):
@@ -129,7 +116,7 @@ def model_fn(model_dir: str) -> dict:
 # Inference (for SageMaker inference)
 # =============================================================================
 def predict_fn(df: pd.DataFrame, model_dict: dict) -> pd.DataFrame:
-    """Make predictions with TabularMLP ensemble."""
+    """Make predictions with PyTorch TabularMLP ensemble."""
     model_type = TEMPLATE_PARAMS["model_type"]
     compressed_features = TEMPLATE_PARAMS["compressed_features"]
     model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
@@ -233,6 +220,36 @@ def predict_fn(df: pd.DataFrame, model_dict: dict) -> pd.DataFrame:
 # Training
 # =============================================================================
 if __name__ == "__main__":
+    # -------------------------------------------------------------------------
+    # Training-only imports (deferred to reduce serverless startup time)
+    # -------------------------------------------------------------------------
+    import argparse
+
+    import awswrangler as wr
+    from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
+    from sklearn.preprocessing import LabelEncoder
+
+    # Enable Tensor Core optimization for GPUs that support it
+    torch.set_float32_matmul_precision("medium")
+
+    from model_script_utils import (
+        check_dataframe,
+        compute_classification_metrics,
+        compute_regression_metrics,
+        print_classification_metrics,
+        print_confusion_matrix,
+        print_regression_metrics,
+    )
+    from pytorch_utils import (
+        create_model,
+        save_model,
+        train_model,
+    )
+    from uq_harness import (
+        save_uq_models,
+        train_uq_models,
+    )
+
     # -------------------------------------------------------------------------
     # Setup: Parse arguments and load data
     # -------------------------------------------------------------------------
@@ -377,6 +394,7 @@ if __name__ == "__main__":
             patience=hyperparameters["early_stopping_patience"],
             batch_size=hyperparameters["batch_size"],
             learning_rate=hyperparameters["learning_rate"],
+            loss=hyperparameters.get("loss", "L1Loss"),
             device=device,
         )
         ensemble_models.append(model)

workbench/model_scripts/pytorch_model/model_script_utils.py

@@ -148,12 +148,16 @@ def convert_categorical_types(
 def decompress_features(
     df: pd.DataFrame, features: list[str], compressed_features: list[str]
 ) -> tuple[pd.DataFrame, list[str]]:
-    """Decompress bitstring features into individual bit columns.
+    """Decompress compressed features (bitstrings or count vectors) into individual columns.
+
+    Supports two formats (auto-detected):
+        - Bitstrings: "10110010..." → individual uint8 columns (0 or 1)
+        - Count vectors: "0,3,0,1,5,..." → individual uint8 columns (0-255)
 
     Args:
         df: The features DataFrame
         features: Full list of feature names
-        compressed_features: List of feature names to decompress (bitstrings)
+        compressed_features: List of feature names to decompress
 
     Returns:
         Tuple of (DataFrame with decompressed features, updated feature list)
@@ -178,18 +182,18 @@ def decompress_features(
         # Remove the feature from the list to avoid duplication
         decompressed_features.remove(feature)
 
-        # Handle all compressed features as bitstrings
-        bit_matrix = np.array([list(bitstring) for bitstring in df[feature]], dtype=np.uint8)
-        prefix = feature[:3]
+        # Auto-detect format and parse: comma-separated counts or bitstring
+        sample = str(df[feature].dropna().iloc[0]) if not df[feature].dropna().empty else ""
+        parse_fn = (lambda s: list(map(int, s.split(",")))) if "," in sample else list
+        feature_matrix = np.array([parse_fn(s) for s in df[feature]], dtype=np.uint8)
 
-        # Create all new columns at once - avoids fragmentation
-        new_col_names = [f"{prefix}_{i}" for i in range(bit_matrix.shape[1])]
-        new_df = pd.DataFrame(bit_matrix, columns=new_col_names, index=df.index)
+        # Create new columns with prefix from feature name
+        prefix = feature[:3]
+        new_col_names = [f"{prefix}_{i}" for i in range(feature_matrix.shape[1])]
+        new_df = pd.DataFrame(feature_matrix, columns=new_col_names, index=df.index)
 
-        # Add to features list
+        # Update features list and dataframe
         decompressed_features.extend(new_col_names)
-
-        # Drop original column and concatenate new ones
         df = df.drop(columns=[feature])
         df = pd.concat([df, new_df], axis=1)