workbench 0.8.207__tar.gz → 0.8.210__tar.gz

{workbench-0.8.207 → workbench-0.8.210}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.207
+Version: 0.8.210
 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
@@ -52,7 +52,7 @@ Requires-Dist: joblib>=1.3.2
 Requires-Dist: requests>=2.26.0
 Requires-Dist: rdkit>=2024.9.5
 Requires-Dist: mordredcommunity>=2.0.6
-Requires-Dist: workbench-bridges>=0.1.15
+Requires-Dist: workbench-bridges>=0.1.16
 Provides-Extra: ui
 Requires-Dist: plotly>=6.0.0; extra == "ui"
 Requires-Dist: dash>3.0.0; extra == "ui"

{workbench-0.8.207 → workbench-0.8.210}/applications/aws_dashboard/Dockerfile

@@ -7,7 +7,7 @@ EXPOSE 8000
 ENV PYTHONHASHSEED=0
 
 # Define workbench version
-ARG WORKBENCH_VERSION=0.8.207
+ARG WORKBENCH_VERSION=0.8.208
 
 # Install system packages and gunicorn in one layer
 RUN apt-get update && apt-get install -y --no-install-recommends \

{workbench-0.8.207 → workbench-0.8.210}/aws_setup/workbench_dashboard_full/app.py

@@ -12,7 +12,7 @@ print(f"Account: {aws_account}")
 print(f"Region: {aws_region}")
 
 # When you want a different docker image change this line
-dashboard_image = "public.ecr.aws/m6i5k1r2/workbench_dashboard:v0_8_207_amd64"
+dashboard_image = "public.ecr.aws/m6i5k1r2/workbench_dashboard:v0_8_208_amd64"
 
 # Workbench Configuration
 try:

{workbench-0.8.207 → workbench-0.8.210}/aws_setup/workbench_dashboard_lite/app.py

@@ -12,7 +12,7 @@ print(f"Account: {aws_account}")
 print(f"Region: {aws_region}")
 
 # When you want a different docker image change this line
-dashboard_image = "public.ecr.aws/m6i5k1r2/workbench_dashboard:v0_8_207_amd64"
+dashboard_image = "public.ecr.aws/m6i5k1r2/workbench_dashboard:v0_8_208_amd64"
 
 # Workbench Configuration
 try:

{workbench-0.8.207 → workbench-0.8.210}/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.196
+RUN pip install --no-cache-dir workbench==0.8.208
 
 # Add the script runner
 COPY ml_pipeline_runner.py /app/ml_pipeline_runner.py

{workbench-0.8.207 → workbench-0.8.210}/pyproject.toml

@@ -41,7 +41,7 @@ dependencies = [
     "requests >= 2.26.0",
     "rdkit>=2024.9.5",
     "mordredcommunity>=2.0.6",
-    "workbench-bridges>=0.1.15",
+    "workbench-bridges>=0.1.16",
 ]
 
 [project.optional-dependencies]

{workbench-0.8.207 → workbench-0.8.210}/src/workbench/core/artifacts/endpoint_core.py

@@ -531,38 +531,14 @@ class EndpointCore(Artifact):
         fs = FeatureSetCore(model.get_input())
         id_column = fs.id_column
 
-        # Is this a UQ Model? If so, run full inference and merge the results
+        # For UQ models, get UQ columns from training CV results and compute metrics
+        # Note: XGBoost training now saves all UQ columns (q_*, confidence, prediction_std)
         additional_columns = []
         if model.model_framework == ModelFramework.XGBOOST and 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 = model.training_view().pull_dataframe()
-
-            # Run inference on the endpoint to get UQ outputs
-            uq_df = self.inference(training_df)
-
-            # Identify UQ-specific columns (quantiles, prediction_std, *_pred_std)
-            uq_columns = [
-                col
-                for col in uq_df.columns
-                if col.startswith("q_") or col == "prediction_std" or col.endswith("_pred_std") or col == "confidence"
-            ]
-
-            # Merge UQ columns with out-of-fold predictions
+            uq_columns = [col for col in out_of_fold_df.columns if col.startswith("q_") or col == "confidence"]
             if uq_columns:
-                # Keep id_column and UQ columns, drop 'prediction' to avoid conflict when merging
-                uq_df = uq_df[[id_column] + uq_columns]
-
-                # Drop duplicates in uq_df based on id_column
-                uq_df = uq_df.drop_duplicates(subset=[id_column])
-
-                # Merge UQ columns into out_of_fold_df
-                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 (use first target for multi-target models)
+                self.log.info(f"UQ columns from training: {', '.join(uq_columns)}")
                 primary_target = targets[0] if isinstance(targets, list) else targets
                 metrics = uq_metrics(out_of_fold_df, primary_target)
                 self.param_store.upsert(f"/workbench/models/{model.name}/inference/full_cross_fold", metrics)

{workbench-0.8.207 → workbench-0.8.210}/src/workbench/model_scripts/chemprop/generated_model_script.py

@@ -58,14 +58,35 @@ from model_script_utils import (
     print_confusion_matrix,
 )
 
+# Default Hyperparameters for ChemProp MPNN
+DEFAULT_HYPERPARAMETERS = {
+    # Training parameters
+    "n_folds": 5,  # Number of CV folds (1 = single train/val split)
+    "max_epochs": 400,  # Maximum training epochs
+    "patience": 40,  # Early stopping patience
+    "batch_size": 16,  # Batch size for training
+
+    # Model architecture parameters (Message Passing)
+    "hidden_dim": 700,  # Hidden dimension for message passing
+    "depth": 6,  # Number of message passing steps
+    "dropout": 0.15,  # Dropout rate
+
+    # FFN parameters
+    "ffn_hidden_dim": 2000,  # Hidden dimension for feed-forward network
+    "ffn_num_layers": 2,  # Number of FFN layers
+
+    # Random seed
+    "seed": 42,
+}
+
 # Template Parameters
 TEMPLATE_PARAMS = {
-    "model_type": "regressor",
-    "targets": ['logd'],  # List of target columns (single or multi-task)
-    "feature_list": ['smiles'],
-    "id_column": "molecule_name",
-    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/open-admet-chemprop-logd/training",
-    "hyperparameters": {'max_epochs': 400, 'hidden_dim': 300, 'depth': 3, 'n_folds': 5},
+    "model_type": "uq_regressor",
+    "targets": ['udm_asy_res_efflux_ratio'],  # List of target columns (single or multi-task)
+    "feature_list": ['smiles', 'smr_vsa4', 'tpsa', 'numhdonors', 'nhohcount', 'nbase', 'vsa_estate3', 'mollogp', 'smr_vsa3', 'peoe_vsa8', 'xc_4d', 'charge_centroid_distance', 'xc_4dv', 'vsa_estate4', 'peoe_vsa1', 'fr_imine', 'estate_vsa9', 'mi', 'molecular_volume_3d', 'fr_guanido', 'slogp_vsa2', 'nitrogen_span', 'slogp_vsa1', 'hba_hbd_ratio', 'asphericity', 'axp_5d', 'axp_0dv', 'estate_vsa4', 'vsa_estate2', 'molecular_axis_length', 'numsaturatedheterocycles', 'xc_3dv', 'minpartialcharge', 'axp_1d', 'chi1v', 'peoe_vsa10', 'fr_imidazole', 'fpdensitymorgan2', 'c3sp3', 'estate_vsa2', 'bcut2d_mrlow', 'kappa3', 'smr_vsa6', 'estate_vsa10', 'minabsestateindex', 'bcut2d_mwhi', 'tertiary_amine_count', 'maxestateindex', 'bcut2d_mrhi', 'molmr', 'qed'],
+    "id_column": "udm_mol_bat_id",
+    "model_metrics_s3_path": "s3://ideaya-sageworks-bucket/models/caco2-er-chemprop-reg-hybrid/training",
+    "hyperparameters": {'n_folds': 5, 'hidden_dim': 700, 'depth': 6, 'dropout': 0.15, 'ffn_hidden_dim': 2000, 'ffn_num_layers': 2},
 }
 
 
@@ -153,12 +174,12 @@ def build_mpnn_model(
     Returns:
         Configured MPNN model
     """
-    # Model hyperparameters with defaults
-    hidden_dim = hyperparameters.get("hidden_dim", 700)
-    depth = hyperparameters.get("depth", 6)
-    dropout = hyperparameters.get("dropout", 0.15)
-    ffn_hidden_dim = hyperparameters.get("ffn_hidden_dim", 2000)
-    ffn_num_layers = hyperparameters.get("ffn_num_layers", 2)
+    # Model hyperparameters (from merged defaults)
+    hidden_dim = hyperparameters["hidden_dim"]
+    depth = hyperparameters["depth"]
+    dropout = hyperparameters["dropout"]
+    ffn_hidden_dim = hyperparameters["ffn_hidden_dim"]
+    ffn_num_layers = hyperparameters["ffn_num_layers"]
 
     # Message passing component
     mp = nn.BondMessagePassing(d_h=hidden_dim, depth=depth, dropout=dropout)
@@ -427,7 +448,7 @@ if __name__ == "__main__":
     feature_list = TEMPLATE_PARAMS["feature_list"]
     id_column = TEMPLATE_PARAMS["id_column"]
     model_metrics_s3_path = TEMPLATE_PARAMS["model_metrics_s3_path"]
-    hyperparameters = TEMPLATE_PARAMS["hyperparameters"] or {}
+    hyperparameters = {**DEFAULT_HYPERPARAMETERS, **(TEMPLATE_PARAMS["hyperparameters"] or {})}
 
     # Validate target_columns
     if not target_columns or not isinstance(target_columns, list) or len(target_columns) == 0:
@@ -505,14 +526,14 @@ if __name__ == "__main__":
     else:
         num_classes = None
 
-    # Model and training configuration
+    # Model and training configuration (from merged defaults)
     print(f"Hyperparameters: {hyperparameters}")
     task = "classification" if model_type == "classifier" else "regression"
     n_extra = len(extra_feature_cols) if use_extra_features else 0
-    max_epochs = hyperparameters.get("max_epochs", 400)
-    patience = hyperparameters.get("patience", 40)
-    n_folds = hyperparameters.get("n_folds", 5)  # Number of CV folds (default: 5)
-    batch_size = hyperparameters.get("batch_size", 16)
+    max_epochs = hyperparameters["max_epochs"]
+    patience = hyperparameters["patience"]
+    n_folds = hyperparameters["n_folds"]
+    batch_size = hyperparameters["batch_size"]
 
     # Check extra feature columns exist
     if use_extra_features:
@@ -639,7 +660,7 @@ if __name__ == "__main__":
         val_loader = data.build_dataloader(val_dataset, batch_size=batch_size, shuffle=False)
 
         # Build and train model for this fold
-        pl.seed_everything(42 + fold_idx)
+        pl.seed_everything(hyperparameters["seed"] + fold_idx)
         mpnn = build_mpnn_model(
             hyperparameters, task=task, num_classes=num_classes, n_targets=n_targets,
             n_extra_descriptors=n_extra, x_d_transform=x_d_transform, output_transform=output_transform,
@@ -841,6 +862,10 @@ if __name__ == "__main__":
     if label_encoder is not None:
         joblib.dump(label_encoder, os.path.join(args.model_dir, "label_encoder.joblib"))
 
+    # Save hyperparameters (merged defaults + user overrides) for reproducibility
+    with open(os.path.join(args.model_dir, "hyperparameters.json"), "w") as fp:
+        json.dump(hyperparameters, fp, indent=2)
+
     # Save extra feature metadata for inference (hybrid mode)
     # Note: We don't need to save the scaler - X_d_transform is embedded in the model
     if use_extra_features:

{workbench-0.8.207 → workbench-0.8.210}/src/workbench/model_scripts/pytorch_model/generated_model_script.py

@@ -2,6 +2,10 @@
 import os
 import awswrangler as wr
 import numpy as np
+import torch
+
+# Enable Tensor Core optimization for GPUs that support it (e.g., NVIDIA L4, A100, etc.)
+torch.set_float32_matmul_precision("medium")
 
 # PyTorch compatibility: pytorch-tabular saves complex objects, not just tensors
 # Use legacy loading behavior for compatibility (recommended by PyTorch docs for this scenario)
@@ -37,15 +41,38 @@ from model_script_utils import (
     print_confusion_matrix,
 )
 
+# Default Hyperparameters for PyTorch Tabular
+DEFAULT_HYPERPARAMETERS = {
+    # Training parameters
+    "n_folds": 5,  # Number of CV folds (1 = single train/val split)
+    "max_epochs": 200,  # Maximum training epochs
+    "min_epochs": 10,  # Minimum epochs before early stopping
+    "early_stopping_patience": 20,  # Patience for early stopping
+    "batch_size": 128,  # Batch size (auto-adjusted if too large for data)
+    "gradient_clip_val": 1.0,  # Gradient clipping value
+    "auto_lr_find": False,  # Whether to auto-find learning rate
+
+    # Model architecture parameters
+    "layers": "256-128-64",  # Hidden layer sizes (dash-separated)
+    "activation": "LeakyReLU",  # Activation function
+    "learning_rate": 1e-3,  # Learning rate
+    "dropout": 0.1,  # Dropout rate
+    "use_batch_norm": True,  # Whether to use batch normalization
+    "initialization": "kaiming",  # Weight initialization method
+
+    # Random seed
+    "seed": 42,
+}
+
 # Template Parameters
 TEMPLATE_PARAMS = {
-    "model_type": "classifier",
-    "target": "solubility_class",
-    "features": ['molwt', 'mollogp', 'molmr', 'heavyatomcount', 'numhacceptors', 'numhdonors', 'numheteroatoms', 'numrotatablebonds', 'numvalenceelectrons', 'numaromaticrings', 'numsaturatedrings', 'numaliphaticrings', 'ringcount', 'tpsa', 'labuteasa', 'balabanj', 'bertzct'],
-    "id_column": "id",
+    "model_type": "uq_regressor",
+    "target": "udm_asy_res_efflux_ratio",
+    "features": ['smr_vsa4', 'tpsa', 'numhdonors', 'nhohcount', 'nbase', 'vsa_estate3', 'mollogp', 'smr_vsa3', 'peoe_vsa8', 'xc_4d', 'charge_centroid_distance', 'xc_4dv', 'vsa_estate4', 'peoe_vsa1', 'fr_imine', 'estate_vsa9', 'mi', 'molecular_volume_3d', 'fr_guanido', 'slogp_vsa2', 'nitrogen_span', 'slogp_vsa1', 'hba_hbd_ratio', 'asphericity', 'axp_5d', 'axp_0dv', 'estate_vsa4', 'vsa_estate2', 'molecular_axis_length', 'numsaturatedheterocycles', 'xc_3dv', 'minpartialcharge', 'axp_1d', 'chi1v', 'peoe_vsa10', 'fr_imidazole', 'fpdensitymorgan2', 'c3sp3', 'estate_vsa2', 'bcut2d_mrlow', 'kappa3', 'smr_vsa6', 'estate_vsa10', 'minabsestateindex', 'bcut2d_mwhi', 'tertiary_amine_count', 'maxestateindex', 'bcut2d_mrhi', 'molmr', 'qed', 'num_s_centers', 'vsa_estate6', 'peoe_vsa2', 'xp_4d', 'hallkieralpha', 'amphiphilic_moment', 'bcut2d_logplow', 'vsa_estate5', 'vsa_estate8', 'xp_1dv', 'vsa_estate7', 'mm', 'smr_vsa5', 'phi', 'amide_count', 'bcut2d_mwlow', 'bcut2d_chglo', 'xp_6d', 'slogp_vsa5', 'chi0n', 'xp_5d', 'axp_2d', 'xpc_5d', 'molecular_asymmetry', 'type_ii_pattern_count', 'smr_vsa1', 'axp_7d', 'xp_2d', 'estate_vsa3', 'xch_7d', 'xp_0dv', 'xc_5dv', 'mp', 'xp_7d', 'slogp_vsa4', 'axp_3d', 'maxabsestateindex', 'kappa2', 'numatomstereocenters', 'axp_7dv', 'chi3n', 'peoe_vsa9', 'axp_1dv', 'axp_3dv', 'maxabspartialcharge', 'vsa_estate9', 'mz', 'peoe_vsa7', 'axp_5dv', 'aromatic_interaction_score', 'axp_0d', 'xp_6dv', 'axp_2dv', 'fpdensitymorgan1', 'radius_of_gyration', 'xpc_4d', 'balabanj', 'xp_3d', 'estate_vsa1', 'sse', 'type_i_pattern_count', 'chi0', 'fr_al_oh', 'fr_ar_n', 'fpdensitymorgan3', 'estate_vsa8', 'xp_7dv', 'axp_6dv', 'slogp_vsa7', 'mv', 'axp_6d', 'peoe_vsa11', 'slogp_vsa3', 'numheterocycles', 'sp', 'estate_vsa6', 'numaromaticheterocycles', 'vsa_estate10', 'xch_4dv', 'chi2v', 'estate_vsa5', 'xp_5dv', 'kappa1', 'peoe_vsa3', 'bcut2d_chghi', 'bcut2d_logphi', 'smr_vsa10', 'axp_4dv', 'xch_7dv', 'chi4v', 'mare', 'xch_6dv', 'axp_4d', 'chi4n', 'smr_vsa9', 'mse', 'peoe_vsa6', 'sv', 'xp_2dv', 'xpc_4dv', 'estate_vsa7', 'smr_vsa7', 'bertzct', 'vsa_estate1', 'si', 'fractioncsp3', 'c1sp3', 'slogp_vsa6', 'sps', 'chi1', 'maxpartialcharge', 'xc_3d', 'minestateindex', 'avgipc', 'xpc_6d', 'xch_5dv', 'slogp_vsa10', 'hybratio', 'xc_5d', 'minabspartialcharge', 'c3sp2', 'mpe', 'fr_methoxy', 'spe', 'xpc_6dv', 'fr_hoccn', 'xch_5d', 'heavyatommolwt', 'c2sp3', 'xch_6d', 'num_r_centers', 'fr_ether', 'xpc_5dv', 'frac_defined_stereo', 'fr_nh0', 'xp_4dv', 'intramolecular_hbond_potential', 'xch_4d', 'xp_1d', 'peoe_vsa13', 'xc_6dv', 'chi1n', 'chi2n', 'xc_6d', 'c1sp2', 'c2sp2', 'num_stereocenters', 'chi3v', 'labuteasa', 'peoe_vsa12', 'numrotatablebonds', 'chi0v', 'numhacceptors', 'exactmolwt', 'fr_piperdine', 'sare', 'fr_aniline', 'numheteroatoms', 'naromatom', 'molwt', 'xch_3d', 'fr_urea', 'peoe_vsa4', 'fr_nh1', 'stereo_complexity', 'slogp_vsa11', 'fr_pyridine', 'xp_3dv', 'fr_amide', 'fr_nitrile', 'fcsp3', 'fr_al_oh_notert', 'fr_halogen', 'nocount', 'fr_ar_nh', 'xp_0d', 'fr_nhpyrrole', 'smr_vsa2', 'fr_ndealkylation2', 'numaliphaticheterocycles', 'numamidebonds', 'fr_aryl_methyl', 'numsaturatedrings', 'sm', 'slogp_vsa8', 'fr_thiazole', 'nacid', 'sz', 'num_defined_stereocenters', 'peoe_vsa5', 'fr_piperzine', 'numaromaticrings', 'fr_benzene', 'fr_oxazole', 'numvalenceelectrons', 'fr_sulfonamd', 'fr_ketone', 'numunspecifiedatomstereocenters'],
+    "id_column": "udm_mol_bat_id",
     "compressed_features": [],
-    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/aqsol-pytorch-class/training",
-    "hyperparameters": {'training_config': {'max_epochs': 150}, 'model_config': {'layers': '256-128-64'}},
+    "model_metrics_s3_path": "s3://ideaya-sageworks-bucket/models/caco2-er-reg-pytorch-test/training",
+    "hyperparameters": {'max_epochs': 200, 'early_stopping_patience': 20, 'layers': '256-128-64', 'dropout': 0.1, 'learning_rate': 0.001, 'activation': 'LeakyReLU'},
 }
 
 
@@ -234,7 +261,7 @@ if __name__ == "__main__":
     compressed_features = TEMPLATE_PARAMS["compressed_features"]
     model_type = TEMPLATE_PARAMS["model_type"]
     model_metrics_s3_path = TEMPLATE_PARAMS["model_metrics_s3_path"]
-    hyperparameters = TEMPLATE_PARAMS["hyperparameters"] or {}
+    hyperparameters = {**DEFAULT_HYPERPARAMETERS, **(TEMPLATE_PARAMS["hyperparameters"] or {})}
 
     # Script arguments for input/output directories
     parser = argparse.ArgumentParser()
@@ -305,9 +332,9 @@ if __name__ == "__main__":
         label_encoder = None
         num_classes = None
 
-    # Use any hyperparameters to set up both the trainer and model configurations
+    # Use hyperparameters (merged with defaults)
     print(f"Hyperparameters: {hyperparameters}")
-    n_folds = hyperparameters.get("n_folds", 5)  # Number of CV folds (default: 5)
+    n_folds = hyperparameters["n_folds"]
 
     # =========================================================================
     # UNIFIED TRAINING: Works for n_folds=1 (single model) or n_folds>1 (K-fold CV)
@@ -352,21 +379,15 @@ if __name__ == "__main__":
         categorical_cols=categorical_cols,
     )
 
-    # Model config defaults
-    model_defaults = {
-        "layers": "256-128-64",
-        "activation": "LeakyReLU",
-        "learning_rate": 1e-3,
-        "dropout": 0.1,
-        "use_batch_norm": True,
-        "initialization": "kaiming",
+    # Model config (from merged hyperparameters)
+    model_params = {
+        "layers": hyperparameters["layers"],
+        "activation": hyperparameters["activation"],
+        "learning_rate": hyperparameters["learning_rate"],
+        "dropout": hyperparameters["dropout"],
+        "use_batch_norm": hyperparameters["use_batch_norm"],
+        "initialization": hyperparameters["initialization"],
     }
-    # Override defaults with model_config if present
-    model_overrides = {k: v for k, v in hyperparameters.get("model_config", {}).items() if k in model_defaults}
-    for key, value in model_overrides.items():
-        print(f"MODEL CONFIG Override: {key}: {model_defaults[key]} → {value}")
-    model_params = {**model_defaults, **model_overrides}
-
     model_config = CategoryEmbeddingModelConfig(task=task, **model_params)
     optimizer_config = OptimizerConfig()
 
@@ -383,30 +404,23 @@ if __name__ == "__main__":
 
         # Set up PyTorch Tabular trainer configuration (per-fold for batch_size)
         # Calculate batch size that avoids single-sample last batch (batch norm requires >1)
-        batch_size = min(128, max(32, len(df_train) // 16))
+        batch_size = min(hyperparameters["batch_size"], max(32, len(df_train) // 16))
         if len(df_train) % batch_size == 1:
             batch_size += 1  # Adjust to avoid last batch of size 1
-        trainer_defaults = {
-            "auto_lr_find": False,
-            "batch_size": batch_size,
-            "max_epochs": 200,
-            "min_epochs": 10,
-            "early_stopping": "valid_loss",
-            "early_stopping_patience": 20,
-            "checkpoints": "valid_loss",
-            "accelerator": "auto",
-            "progress_bar": "none",
-            "gradient_clip_val": 1.0,
-            "seed": 42 + fold_idx,
-        }
-
-        # Override defaults with training_config if present
-        training_overrides = {k: v for k, v in hyperparameters.get("training_config", {}).items() if k in trainer_defaults}
-        if fold_idx == 0:  # Only print overrides once
-            for key, value in training_overrides.items():
-                print(f"TRAINING CONFIG Override: {key}: {trainer_defaults[key]} → {value}")
-        trainer_params = {**trainer_defaults, **training_overrides}
-        trainer_config = TrainerConfig(**trainer_params)
+
+        trainer_config = TrainerConfig(
+            auto_lr_find=hyperparameters["auto_lr_find"],
+            batch_size=batch_size,
+            max_epochs=hyperparameters["max_epochs"],
+            min_epochs=hyperparameters["min_epochs"],
+            early_stopping="valid_loss",
+            early_stopping_patience=hyperparameters["early_stopping_patience"],
+            checkpoints="valid_loss",
+            accelerator="auto",
+            progress_bar="none",
+            gradient_clip_val=hyperparameters["gradient_clip_val"],
+            seed=hyperparameters["seed"] + fold_idx,
+        )
 
         # Create and train the TabularModel for this fold
         tabular_model = TabularModel(
@@ -536,3 +550,7 @@ if __name__ == "__main__":
     # Save the category mappings
     with open(os.path.join(args.model_dir, "category_mappings.json"), "w") as fp:
         json.dump(category_mappings, fp)
+
+    # Save hyperparameters (merged defaults + user overrides) for reproducibility
+    with open(os.path.join(args.model_dir, "hyperparameters.json"), "w") as fp:
+        json.dump(hyperparameters, fp, indent=2)