workbench 0.8.212__py3-none-any.whl → 0.8.217__py3-none-any.whl

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,
 }
@@ -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/pytorch.template

@@ -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,
 }
@@ -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/pytorch_utils.py

@@ -245,6 +245,7 @@ def train_model(
     patience: int = 20,
     batch_size: int = 128,
     learning_rate: float = 1e-3,
+    loss: str = "L1Loss",
     device: str = "cpu",
 ) -> tuple[TabularMLP, dict]:
     """Train the model with early stopping.
@@ -272,7 +273,16 @@ def train_model(
     if task == "classification":
         criterion = nn.CrossEntropyLoss()
     else:
-        criterion = nn.MSELoss()
+        # Map loss name to PyTorch loss class
+        loss_map = {
+            "L1Loss": nn.L1Loss,
+            "MSELoss": nn.MSELoss,
+            "HuberLoss": nn.HuberLoss,
+            "SmoothL1Loss": nn.SmoothL1Loss,
+        }
+        if loss not in loss_map:
+            raise ValueError(f"Unknown loss '{loss}'. Supported: {list(loss_map.keys())}")
+        criterion = loss_map[loss]()
 
     optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
 

workbench/model_scripts/script_generation.py

@@ -100,6 +100,7 @@ def generate_model_script(template_params: dict) -> str:
             - model_metrics_s3_path (str): The S3 path to store the model metrics
             - train_all_data (bool): Whether to train on all (100%) of the data
             - hyperparameters (dict, optional): Hyperparameters for the model (default: None)
+            - child_endpoints (list[str], optional): For META models, list of child endpoint names
 
     Returns:
         str: The name of the generated model script
@@ -116,6 +117,9 @@ def generate_model_script(template_params: dict) -> str:
     elif template_params["model_framework"] == ModelFramework.CHEMPROP:
         template_name = "chemprop.template"
         model_script_dir = "chemprop"
+    elif template_params["model_framework"] == ModelFramework.META:
+        template_name = "meta_model.template"
+        model_script_dir = "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"