workbench 0.8.162__py3-none-any.whl → 0.8.220__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of workbench might be problematic. Click here for more details.

Files changed (147) hide show
  1. workbench/algorithms/dataframe/__init__.py +1 -2
  2. workbench/algorithms/dataframe/compound_dataset_overlap.py +321 -0
  3. workbench/algorithms/dataframe/feature_space_proximity.py +168 -75
  4. workbench/algorithms/dataframe/fingerprint_proximity.py +422 -86
  5. workbench/algorithms/dataframe/projection_2d.py +44 -21
  6. workbench/algorithms/dataframe/proximity.py +259 -305
  7. workbench/algorithms/graph/light/proximity_graph.py +14 -12
  8. workbench/algorithms/models/cleanlab_model.py +382 -0
  9. workbench/algorithms/models/noise_model.py +388 -0
  10. workbench/algorithms/sql/outliers.py +3 -3
  11. workbench/api/__init__.py +5 -1
  12. workbench/api/compound.py +1 -1
  13. workbench/api/df_store.py +17 -108
  14. workbench/api/endpoint.py +18 -5
  15. workbench/api/feature_set.py +121 -15
  16. workbench/api/meta.py +5 -2
  17. workbench/api/meta_model.py +289 -0
  18. workbench/api/model.py +55 -21
  19. workbench/api/monitor.py +1 -16
  20. workbench/api/parameter_store.py +3 -52
  21. workbench/cached/cached_model.py +4 -4
  22. workbench/core/artifacts/__init__.py +11 -2
  23. workbench/core/artifacts/artifact.py +16 -8
  24. workbench/core/artifacts/data_capture_core.py +355 -0
  25. workbench/core/artifacts/df_store_core.py +114 -0
  26. workbench/core/artifacts/endpoint_core.py +382 -253
  27. workbench/core/artifacts/feature_set_core.py +249 -45
  28. workbench/core/artifacts/model_core.py +135 -80
  29. workbench/core/artifacts/monitor_core.py +33 -248
  30. workbench/core/artifacts/parameter_store_core.py +98 -0
  31. workbench/core/cloud_platform/aws/aws_account_clamp.py +50 -1
  32. workbench/core/cloud_platform/aws/aws_meta.py +12 -5
  33. workbench/core/cloud_platform/aws/aws_session.py +4 -4
  34. workbench/core/pipelines/pipeline_executor.py +1 -1
  35. workbench/core/transforms/data_to_features/light/molecular_descriptors.py +4 -4
  36. workbench/core/transforms/features_to_model/features_to_model.py +62 -40
  37. workbench/core/transforms/model_to_endpoint/model_to_endpoint.py +76 -15
  38. workbench/core/transforms/pandas_transforms/pandas_to_features.py +38 -2
  39. workbench/core/views/training_view.py +113 -42
  40. workbench/core/views/view.py +53 -3
  41. workbench/core/views/view_utils.py +4 -4
  42. workbench/model_script_utils/model_script_utils.py +339 -0
  43. workbench/model_script_utils/pytorch_utils.py +405 -0
  44. workbench/model_script_utils/uq_harness.py +278 -0
  45. workbench/model_scripts/chemprop/chemprop.template +649 -0
  46. workbench/model_scripts/chemprop/generated_model_script.py +649 -0
  47. workbench/model_scripts/chemprop/model_script_utils.py +339 -0
  48. workbench/model_scripts/chemprop/requirements.txt +3 -0
  49. workbench/model_scripts/custom_models/chem_info/fingerprints.py +175 -0
  50. workbench/model_scripts/custom_models/chem_info/mol_descriptors.py +483 -0
  51. workbench/model_scripts/custom_models/chem_info/mol_standardize.py +450 -0
  52. workbench/model_scripts/custom_models/chem_info/molecular_descriptors.py +7 -9
  53. workbench/model_scripts/custom_models/chem_info/morgan_fingerprints.py +1 -1
  54. workbench/model_scripts/custom_models/proximity/feature_space_proximity.py +194 -0
  55. workbench/model_scripts/custom_models/proximity/feature_space_proximity.template +8 -10
  56. workbench/model_scripts/custom_models/uq_models/bayesian_ridge.template +7 -8
  57. workbench/model_scripts/custom_models/uq_models/ensemble_xgb.template +20 -21
  58. workbench/model_scripts/custom_models/uq_models/feature_space_proximity.py +194 -0
  59. workbench/model_scripts/custom_models/uq_models/gaussian_process.template +5 -11
  60. workbench/model_scripts/custom_models/uq_models/ngboost.template +30 -18
  61. workbench/model_scripts/custom_models/uq_models/requirements.txt +1 -3
  62. workbench/model_scripts/ensemble_xgb/ensemble_xgb.template +15 -17
  63. workbench/model_scripts/meta_model/generated_model_script.py +209 -0
  64. workbench/model_scripts/meta_model/meta_model.template +209 -0
  65. workbench/model_scripts/pytorch_model/generated_model_script.py +444 -500
  66. workbench/model_scripts/pytorch_model/model_script_utils.py +339 -0
  67. workbench/model_scripts/pytorch_model/pytorch.template +440 -496
  68. workbench/model_scripts/pytorch_model/pytorch_utils.py +405 -0
  69. workbench/model_scripts/pytorch_model/requirements.txt +1 -1
  70. workbench/model_scripts/pytorch_model/uq_harness.py +278 -0
  71. workbench/model_scripts/scikit_learn/generated_model_script.py +7 -12
  72. workbench/model_scripts/scikit_learn/scikit_learn.template +4 -9
  73. workbench/model_scripts/script_generation.py +20 -11
  74. workbench/model_scripts/uq_models/generated_model_script.py +248 -0
  75. workbench/model_scripts/xgb_model/generated_model_script.py +372 -404
  76. workbench/model_scripts/xgb_model/model_script_utils.py +339 -0
  77. workbench/model_scripts/xgb_model/uq_harness.py +278 -0
  78. workbench/model_scripts/xgb_model/xgb_model.template +369 -401
  79. workbench/repl/workbench_shell.py +28 -19
  80. workbench/resources/open_source_api.key +1 -1
  81. workbench/scripts/endpoint_test.py +162 -0
  82. workbench/scripts/lambda_test.py +73 -0
  83. workbench/scripts/meta_model_sim.py +35 -0
  84. workbench/scripts/ml_pipeline_batch.py +137 -0
  85. workbench/scripts/ml_pipeline_sqs.py +186 -0
  86. workbench/scripts/monitor_cloud_watch.py +20 -100
  87. workbench/scripts/training_test.py +85 -0
  88. workbench/utils/aws_utils.py +4 -3
  89. workbench/utils/chem_utils/__init__.py +0 -0
  90. workbench/utils/chem_utils/fingerprints.py +175 -0
  91. workbench/utils/chem_utils/misc.py +194 -0
  92. workbench/utils/chem_utils/mol_descriptors.py +483 -0
  93. workbench/utils/chem_utils/mol_standardize.py +450 -0
  94. workbench/utils/chem_utils/mol_tagging.py +348 -0
  95. workbench/utils/chem_utils/projections.py +219 -0
  96. workbench/utils/chem_utils/salts.py +256 -0
  97. workbench/utils/chem_utils/sdf.py +292 -0
  98. workbench/utils/chem_utils/toxicity.py +250 -0
  99. workbench/utils/chem_utils/vis.py +253 -0
  100. workbench/utils/chemprop_utils.py +141 -0
  101. workbench/utils/cloudwatch_handler.py +1 -1
  102. workbench/utils/cloudwatch_utils.py +137 -0
  103. workbench/utils/config_manager.py +3 -7
  104. workbench/utils/endpoint_utils.py +5 -7
  105. workbench/utils/license_manager.py +2 -6
  106. workbench/utils/meta_model_simulator.py +499 -0
  107. workbench/utils/metrics_utils.py +256 -0
  108. workbench/utils/model_utils.py +278 -79
  109. workbench/utils/monitor_utils.py +44 -62
  110. workbench/utils/pandas_utils.py +3 -3
  111. workbench/utils/pytorch_utils.py +87 -0
  112. workbench/utils/shap_utils.py +11 -57
  113. workbench/utils/workbench_logging.py +0 -3
  114. workbench/utils/workbench_sqs.py +1 -1
  115. workbench/utils/xgboost_local_crossfold.py +267 -0
  116. workbench/utils/xgboost_model_utils.py +127 -219
  117. workbench/web_interface/components/model_plot.py +14 -2
  118. workbench/web_interface/components/plugin_unit_test.py +5 -2
  119. workbench/web_interface/components/plugins/dashboard_status.py +3 -1
  120. workbench/web_interface/components/plugins/generated_compounds.py +1 -1
  121. workbench/web_interface/components/plugins/model_details.py +38 -74
  122. workbench/web_interface/components/plugins/scatter_plot.py +6 -10
  123. {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/METADATA +31 -9
  124. {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/RECORD +128 -96
  125. workbench-0.8.220.dist-info/entry_points.txt +11 -0
  126. {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/licenses/LICENSE +1 -1
  127. workbench/core/cloud_platform/aws/aws_df_store.py +0 -404
  128. workbench/core/cloud_platform/aws/aws_parameter_store.py +0 -280
  129. workbench/model_scripts/custom_models/chem_info/local_utils.py +0 -769
  130. workbench/model_scripts/custom_models/chem_info/tautomerize.py +0 -83
  131. workbench/model_scripts/custom_models/meta_endpoints/example.py +0 -53
  132. workbench/model_scripts/custom_models/proximity/generated_model_script.py +0 -138
  133. workbench/model_scripts/custom_models/proximity/proximity.py +0 -384
  134. workbench/model_scripts/custom_models/uq_models/generated_model_script.py +0 -393
  135. workbench/model_scripts/custom_models/uq_models/mapie_xgb.template +0 -203
  136. workbench/model_scripts/custom_models/uq_models/meta_uq.template +0 -273
  137. workbench/model_scripts/custom_models/uq_models/proximity.py +0 -384
  138. workbench/model_scripts/ensemble_xgb/generated_model_script.py +0 -279
  139. workbench/model_scripts/quant_regression/quant_regression.template +0 -279
  140. workbench/model_scripts/quant_regression/requirements.txt +0 -1
  141. workbench/utils/chem_utils.py +0 -1556
  142. workbench/utils/execution_environment.py +0 -211
  143. workbench/utils/fast_inference.py +0 -167
  144. workbench/utils/resource_utils.py +0 -39
  145. workbench-0.8.162.dist-info/entry_points.txt +0 -5
  146. {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/WHEEL +0 -0
  147. {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/top_level.txt +0 -0
@@ -1,393 +0,0 @@
1
- # Model: NGBoost Regressor with Distribution output
2
- from ngboost import NGBRegressor
3
- from xgboost import XGBRegressor # Base Estimator
4
- from sklearn.model_selection import train_test_split
5
- import numpy as np
6
-
7
- # Model Performance Scores
8
- from sklearn.metrics import (
9
- mean_absolute_error,
10
- r2_score,
11
- root_mean_squared_error
12
- )
13
-
14
- from io import StringIO
15
- import json
16
- import argparse
17
- import joblib
18
- import os
19
- import pandas as pd
20
-
21
- # Local Imports
22
- from proximity import Proximity
23
-
24
-
25
-
26
- # Template Placeholders
27
- TEMPLATE_PARAMS = {
28
- "id_column": "id",
29
- "features": ['molwt', 'mollogp', 'molmr', 'heavyatomcount', 'numhacceptors', 'numhdonors', 'numheteroatoms', 'numrotatablebonds', 'numvalenceelectrons', 'numaromaticrings', 'numsaturatedrings', 'numaliphaticrings', 'ringcount', 'tpsa', 'labuteasa', 'balabanj', 'bertzct'],
30
- "target": "solubility",
31
- "train_all_data": True,
32
- "track_columns": ['solubility']
33
- }
34
-
35
-
36
- # Function to check if dataframe is empty
37
- def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
38
- """
39
- Check if the provided dataframe is empty and raise an exception if it is.
40
-
41
- Args:
42
- df (pd.DataFrame): DataFrame to check
43
- df_name (str): Name of the DataFrame
44
- """
45
- if df.empty:
46
- msg = f"*** The training data {df_name} has 0 rows! ***STOPPING***"
47
- print(msg)
48
- raise ValueError(msg)
49
-
50
-
51
- def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
52
- """
53
- Matches and renames DataFrame columns to match model feature names (case-insensitive).
54
- Prioritizes exact matches, then case-insensitive matches.
55
-
56
- Raises ValueError if any model features cannot be matched.
57
- """
58
- df_columns_lower = {col.lower(): col for col in df.columns}
59
- rename_dict = {}
60
- missing = []
61
- for feature in model_features:
62
- if feature in df.columns:
63
- continue # Exact match
64
- elif feature.lower() in df_columns_lower:
65
- rename_dict[df_columns_lower[feature.lower()]] = feature
66
- else:
67
- missing.append(feature)
68
-
69
- if missing:
70
- raise ValueError(f"Features not found: {missing}")
71
-
72
- # Rename the DataFrame columns to match the model features
73
- return df.rename(columns=rename_dict)
74
-
75
-
76
- def distance_weighted_calibrated_intervals(
77
- df_pred: pd.DataFrame,
78
- prox_df: pd.DataFrame,
79
- calibration_strength: float = 0.7,
80
- distance_decay: float = 3.0,
81
- ) -> pd.DataFrame:
82
- """
83
- Calibrate intervals using distance-weighted neighbor quantiles.
84
- Uses all 10 neighbors with distance-based weighting.
85
- """
86
- id_column = TEMPLATE_PARAMS["id_column"]
87
- target_column = TEMPLATE_PARAMS["target"]
88
-
89
- # Distance-weighted neighbor statistics
90
- def weighted_quantile(values, weights, q):
91
- """Calculate weighted quantile"""
92
- if len(values) == 0:
93
- return np.nan
94
- sorted_indices = np.argsort(values)
95
- sorted_values = values[sorted_indices]
96
- sorted_weights = weights[sorted_indices]
97
- cumsum = np.cumsum(sorted_weights)
98
- cutoff = q * cumsum[-1]
99
- return np.interp(cutoff, cumsum, sorted_values)
100
-
101
- # Calculate distance weights (closer neighbors get more weight)
102
- prox_df = prox_df.copy()
103
- prox_df['weight'] = 1 / (1 + prox_df['distance'] ** distance_decay)
104
-
105
- # Get weighted quantiles and statistics for each ID
106
- neighbor_stats = []
107
- for id_val, group in prox_df.groupby(id_column):
108
- values = group[target_column].values
109
- weights = group['weight'].values
110
-
111
- # Normalize weights
112
- weights = weights / weights.sum()
113
-
114
- stats = {
115
- id_column: id_val,
116
- 'local_q025': weighted_quantile(values, weights, 0.025),
117
- 'local_q25': weighted_quantile(values, weights, 0.25),
118
- 'local_q75': weighted_quantile(values, weights, 0.75),
119
- 'local_q975': weighted_quantile(values, weights, 0.975),
120
- 'local_median': weighted_quantile(values, weights, 0.5),
121
- 'local_std': np.sqrt(np.average((values - np.average(values, weights=weights)) ** 2, weights=weights)),
122
- 'avg_distance': group['distance'].mean(),
123
- 'min_distance': group['distance'].min(),
124
- 'max_distance': group['distance'].max(),
125
- }
126
- neighbor_stats.append(stats)
127
-
128
- neighbor_df = pd.DataFrame(neighbor_stats)
129
- out = df_pred.merge(neighbor_df, on=id_column, how='left')
130
-
131
- # Model disagreement score (normalized by prediction std)
132
- model_disagreement = (out["prediction"] - out["prediction_uq"]).abs()
133
- disagreement_score = (model_disagreement / out["prediction_std"]).clip(0, 2)
134
-
135
- # Local confidence based on:
136
- # 1. How close the neighbors are (closer = more confident)
137
- # 2. How much local variance there is (less variance = more confident)
138
- max_reasonable_distance = out['max_distance'].quantile(0.8) # 80th percentile as reference
139
- distance_confidence = (1 - (out['avg_distance'] / max_reasonable_distance)).clip(0.1, 1.0)
140
-
141
- variance_confidence = (out["prediction_std"] / out["local_std"]).clip(0.5, 2.0)
142
- local_confidence = distance_confidence * variance_confidence.clip(0.5, 1.5)
143
-
144
- # Calibration weight: higher when models disagree and we have good local data
145
- calibration_weight = (
146
- calibration_strength *
147
- local_confidence * # Weight by local data quality
148
- disagreement_score.clip(0.3, 1.0) # More calibration when models disagree
149
- )
150
-
151
- # Consensus prediction (slight preference for NGBoost since it provides intervals)
152
- consensus_pred = 0.65 * out["prediction_uq"] + 0.35 * out["prediction"]
153
-
154
- # Re-center local intervals around consensus prediction
155
- local_center_offset = consensus_pred - out["local_median"]
156
-
157
- # Apply calibration to each quantile
158
- quantile_pairs = [
159
- ("q_025", "local_q025"),
160
- ("q_25", "local_q25"),
161
- ("q_75", "local_q75"),
162
- ("q_975", "local_q975")
163
- ]
164
-
165
- for model_q, local_q in quantile_pairs:
166
- # Adjust local quantiles to be centered around consensus
167
- adjusted_local_q = out[local_q] + local_center_offset
168
-
169
- # Blend model and local intervals
170
- out[model_q] = (
171
- (1 - calibration_weight) * out[model_q] +
172
- calibration_weight * adjusted_local_q
173
- )
174
-
175
- # Ensure proper interval ordering and bounds using pandas
176
- out["q_025"] = pd.concat([out["q_025"], consensus_pred], axis=1).min(axis=1)
177
- out["q_975"] = pd.concat([out["q_975"], consensus_pred], axis=1).max(axis=1)
178
- out["q_25"] = pd.concat([out["q_25"], out["q_75"]], axis=1).min(axis=1)
179
-
180
- # Optional: Add some interval expansion when neighbors are very far
181
- # (indicates we're in a sparse region of feature space)
182
- sparse_region_mask = out['min_distance'] > out['min_distance'].quantile(0.9)
183
- expansion_factor = 1 + 0.2 * sparse_region_mask # 20% expansion in sparse regions
184
-
185
- for q in ["q_025", "q_25", "q_75", "q_975"]:
186
- interval_width = out[q] - consensus_pred
187
- out[q] = consensus_pred + interval_width * expansion_factor
188
-
189
- # Clean up temporary columns
190
- cleanup_cols = [col for col in out.columns if col.startswith("local_")] + \
191
- ['avg_distance', 'min_distance', 'max_distance']
192
-
193
- return out.drop(columns=cleanup_cols)
194
-
195
-
196
- # TRAINING SECTION
197
- #
198
- # This section (__main__) is where SageMaker will execute the training job
199
- # and save the model artifacts to the model directory.
200
- #
201
- if __name__ == "__main__":
202
- # Template Parameters
203
- id_column = TEMPLATE_PARAMS["id_column"]
204
- features = TEMPLATE_PARAMS["features"]
205
- target = TEMPLATE_PARAMS["target"]
206
- train_all_data = TEMPLATE_PARAMS["train_all_data"]
207
- track_columns = TEMPLATE_PARAMS["track_columns"] # Can be None
208
- validation_split = 0.2
209
-
210
- # Script arguments for input/output directories
211
- parser = argparse.ArgumentParser()
212
- parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
213
- parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
214
- parser.add_argument(
215
- "--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data")
216
- )
217
- args = parser.parse_args()
218
-
219
- # Load training data from the specified directory
220
- training_files = [
221
- os.path.join(args.train, file)
222
- for file in os.listdir(args.train) if file.endswith(".csv")
223
- ]
224
- print(f"Training Files: {training_files}")
225
-
226
- # Combine files and read them all into a single pandas dataframe
227
- df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
228
-
229
- # Check if the DataFrame is empty
230
- check_dataframe(df, "training_df")
231
-
232
- # Training data split logic
233
- if train_all_data:
234
- # Use all data for both training and validation
235
- print("Training on all data...")
236
- df_train = df.copy()
237
- df_val = df.copy()
238
- elif "training" in df.columns:
239
- # Split data based on a 'training' column if it exists
240
- print("Splitting data based on 'training' column...")
241
- df_train = df[df["training"]].copy()
242
- df_val = df[~df["training"]].copy()
243
- else:
244
- # Perform a random split if no 'training' column is found
245
- print("Splitting data randomly...")
246
- df_train, df_val = train_test_split(df, test_size=validation_split, random_state=42)
247
-
248
- # We're using XGBoost for point predictions and NGBoost for uncertainty quantification
249
- xgb_model = XGBRegressor()
250
- ngb_model = NGBRegressor()
251
-
252
- # Prepare features and targets for training
253
- X_train = df_train[features]
254
- X_val = df_val[features]
255
- y_train = df_train[target]
256
- y_val = df_val[target]
257
-
258
- # Train both models using the training data
259
- xgb_model.fit(X_train, y_train)
260
- ngb_model.fit(X_train, y_train, X_val=X_val, Y_val=y_val)
261
-
262
- # Make Predictions on the Validation Set
263
- print(f"Making Predictions on Validation Set...")
264
- y_validate = df_val[target]
265
- X_validate = df_val[features]
266
- preds = xgb_model.predict(X_validate)
267
-
268
- # Calculate various model performance metrics (regression)
269
- rmse = root_mean_squared_error(y_validate, preds)
270
- mae = mean_absolute_error(y_validate, preds)
271
- r2 = r2_score(y_validate, preds)
272
- print(f"RMSE: {rmse:.3f}")
273
- print(f"MAE: {mae:.3f}")
274
- print(f"R2: {r2:.3f}")
275
- print(f"NumRows: {len(df_val)}")
276
-
277
- # Save the trained XGBoost model
278
- xgb_model.save_model(os.path.join(args.model_dir, "xgb_model.json"))
279
-
280
- # Save the trained NGBoost model
281
- joblib.dump(ngb_model, os.path.join(args.model_dir, "ngb_model.joblib"))
282
-
283
- # Save the feature list to validate input during predictions
284
- with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
285
- json.dump(features, fp)
286
-
287
- # Now the Proximity model
288
- model = Proximity(df_train, id_column, features, target, track_columns=track_columns)
289
-
290
- # Now serialize the model
291
- model.serialize(args.model_dir)
292
-
293
-
294
- #
295
- # Inference Section
296
- #
297
- def model_fn(model_dir) -> dict:
298
- """Load and return XGBoost and NGBoost regressors from model directory."""
299
-
300
- # Load XGBoost regressor
301
- xgb_path = os.path.join(model_dir, "xgb_model.json")
302
- xgb_model = XGBRegressor(enable_categorical=True)
303
- xgb_model.load_model(xgb_path)
304
-
305
- # Load NGBoost regressor
306
- ngb_model = joblib.load(os.path.join(model_dir, "ngb_model.joblib"))
307
-
308
- # Deserialize the proximity model
309
- prox_model = Proximity.deserialize(model_dir)
310
-
311
- return {
312
- "xgboost": xgb_model,
313
- "ngboost": ngb_model,
314
- "proximity": prox_model
315
- }
316
-
317
-
318
- def input_fn(input_data, content_type):
319
- """Parse input data and return a DataFrame."""
320
- if not input_data:
321
- raise ValueError("Empty input data is not supported!")
322
-
323
- # Decode bytes to string if necessary
324
- if isinstance(input_data, bytes):
325
- input_data = input_data.decode("utf-8")
326
-
327
- if "text/csv" in content_type:
328
- return pd.read_csv(StringIO(input_data))
329
- elif "application/json" in content_type:
330
- return pd.DataFrame(json.loads(input_data)) # Assumes JSON array of records
331
- else:
332
- raise ValueError(f"{content_type} not supported!")
333
-
334
-
335
- def output_fn(output_df, accept_type):
336
- """Supports both CSV and JSON output formats."""
337
- if "text/csv" in accept_type:
338
- csv_output = output_df.fillna("N/A").to_csv(index=False) # CSV with N/A for missing values
339
- return csv_output, "text/csv"
340
- elif "application/json" in accept_type:
341
- return output_df.to_json(orient="records"), "application/json" # JSON array of records (NaNs -> null)
342
- else:
343
- raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
344
-
345
-
346
- def predict_fn(df, models) -> pd.DataFrame:
347
- """Make Predictions with our XGB Quantile Regression Model
348
-
349
- Args:
350
- df (pd.DataFrame): The input DataFrame
351
- models (dict): The dictionary of models to use for predictions
352
-
353
- Returns:
354
- pd.DataFrame: The DataFrame with the predictions added
355
- """
356
-
357
- # Grab our feature columns (from training)
358
- model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
359
- with open(os.path.join(model_dir, "feature_columns.json")) as fp:
360
- model_features = json.load(fp)
361
-
362
- # Match features in a case-insensitive manner
363
- matched_df = match_features_case_insensitive(df, model_features)
364
-
365
- # Use XGBoost for point predictions
366
- df["prediction"] = models["xgboost"].predict(matched_df[model_features])
367
-
368
- # NGBoost predict returns distribution objects
369
- y_dists = models["ngboost"].pred_dist(matched_df[model_features])
370
-
371
- # Extract parameters from distribution
372
- dist_params = y_dists.params
373
-
374
- # Extract mean and std from distribution parameters
375
- df["prediction_uq"] = dist_params['loc'] # mean
376
- df["prediction_std"] = dist_params['scale'] # standard deviation
377
-
378
- # Add 95% prediction intervals using ppf (percent point function)
379
- df["q_025"] = y_dists.ppf(0.025) # 2.5th percentile
380
- df["q_975"] = y_dists.ppf(0.975) # 97.5th percentile
381
-
382
- # Add 50% prediction intervals
383
- df["q_25"] = y_dists.ppf(0.25) # 25th percentile
384
- df["q_75"] = y_dists.ppf(0.75) # 75th percentile
385
-
386
- # Compute Nearest neighbors with Proximity model
387
- prox_df = models["proximity"].neighbors(df)
388
-
389
- # Shrink prediction intervals based on KNN variance
390
- df = distance_weighted_calibrated_intervals(df, prox_df)
391
-
392
- # Return the modified DataFrame
393
- return df
@@ -1,203 +0,0 @@
1
- # Model: HistGradientBoosting with MAPIE Conformalized Quantile Regression
2
- from mapie.regression import MapieQuantileRegressor
3
- from sklearn.ensemble import HistGradientBoostingRegressor
4
- from sklearn.model_selection import train_test_split
5
- import numpy as np
6
-
7
- # Template Placeholders
8
- TEMPLATE_PARAMS = {
9
- "features": "{{feature_list}}",
10
- "target": "{{target_column}}",
11
- "train_all_data": "{{train_all_data}}"
12
- }
13
-
14
- from io import StringIO
15
- import json
16
- import argparse
17
- import joblib
18
- import os
19
- import pandas as pd
20
-
21
-
22
- # Function to check if dataframe is empty
23
- def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
24
- """Check if the DataFrame is empty and raise an error if so."""
25
- if df.empty:
26
- msg = f"*** The training data {df_name} has 0 rows! ***STOPPING***"
27
- print(msg)
28
- raise ValueError(msg)
29
-
30
-
31
- def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
32
- """
33
- Matches and renames DataFrame columns to match model feature names (case-insensitive).
34
- Prioritizes exact matches, then case-insensitive matches.
35
-
36
- Raises ValueError if any model features cannot be matched.
37
- """
38
- df_columns_lower = {col.lower(): col for col in df.columns}
39
- rename_dict = {}
40
- missing = []
41
- for feature in model_features:
42
- if feature in df.columns:
43
- continue # Exact match
44
- elif feature.lower() in df_columns_lower:
45
- rename_dict[df_columns_lower[feature.lower()]] = feature
46
- else:
47
- missing.append(feature)
48
-
49
- if missing:
50
- raise ValueError(f"Features not found: {missing}")
51
-
52
- # Rename the DataFrame columns to match the model features
53
- return df.rename(columns=rename_dict)
54
-
55
-
56
- # TRAINING SECTION
57
- #
58
- # This section (__main__) is where SageMaker will execute the training job
59
- # and save the model artifacts to the model directory.
60
- #
61
- if __name__ == "__main__":
62
- # Template Parameters
63
- features = TEMPLATE_PARAMS["features"]
64
- target = TEMPLATE_PARAMS["target"]
65
- train_all_data = TEMPLATE_PARAMS["train_all_data"]
66
- validation_split = 0.2
67
-
68
- # Script arguments for input/output directories
69
- parser = argparse.ArgumentParser()
70
- parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
71
- parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
72
- parser.add_argument(
73
- "--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data")
74
- )
75
- args = parser.parse_args()
76
-
77
- # Load training data from the specified directory
78
- training_files = [
79
- os.path.join(args.train, file)
80
- for file in os.listdir(args.train) if file.endswith(".csv")
81
- ]
82
- df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
83
-
84
- # Check if the DataFrame is empty
85
- check_dataframe(df, "training_df")
86
-
87
- # Training data split logic
88
- if train_all_data:
89
- # Use all data for both training and validation
90
- print("Training on all data...")
91
- df_train = df.copy()
92
- df_val = df.copy()
93
- elif "training" in df.columns:
94
- # Split data based on a 'training' column if it exists
95
- print("Splitting data based on 'training' column...")
96
- df_train = df[df["training"]].copy()
97
- df_val = df[~df["training"]].copy()
98
- else:
99
- # Perform a random split if no 'training' column is found
100
- print("Splitting data randomly...")
101
- df_train, df_val = train_test_split(df, test_size=validation_split, random_state=42)
102
-
103
- # Create HistGradientBoosting base model configured for quantile regression
104
- base_estimator = HistGradientBoostingRegressor(
105
- loss='quantile', # Required for MAPIE CQR
106
- quantile=0.5, # Will be overridden by MAPIE for different quantiles
107
- max_iter=1000,
108
- max_depth=6,
109
- learning_rate=0.01,
110
- random_state=42
111
- )
112
-
113
- # Create MAPIE CQR predictor - it will create quantile versions internally
114
- model = MapieQuantileRegressor(
115
- estimator=base_estimator,
116
- method="quantile",
117
- cv="split",
118
- alpha=0.05 # For 95% coverage
119
- )
120
-
121
- # Prepare features and targets for training
122
- X_train = df_train[features]
123
- X_val = df_val[features]
124
- y_train = df_train[target]
125
- y_val = df_val[target]
126
-
127
- # Fit the MAPIE CQR model (train/calibration is handled internally)
128
- model.fit(X_train, y_train)
129
-
130
- # Save the trained model and any necessary assets
131
- joblib.dump(model, os.path.join(args.model_dir, "model.joblib"))
132
-
133
- # Save the feature list to validate input during predictions
134
- with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
135
- json.dump(features, fp)
136
-
137
-
138
- #
139
- # Inference Section
140
- #
141
- def model_fn(model_dir):
142
- """Load and return the model from the specified directory."""
143
- return joblib.load(os.path.join(model_dir, "model.joblib"))
144
-
145
-
146
- def input_fn(input_data, content_type):
147
- """Parse input data and return a DataFrame."""
148
- if not input_data:
149
- raise ValueError("Empty input data is not supported!")
150
-
151
- # Decode bytes to string if necessary
152
- if isinstance(input_data, bytes):
153
- input_data = input_data.decode("utf-8")
154
-
155
- if "text/csv" in content_type:
156
- return pd.read_csv(StringIO(input_data))
157
- elif "application/json" in content_type:
158
- return pd.DataFrame(json.loads(input_data)) # Assumes JSON array of records
159
- else:
160
- raise ValueError(f"{content_type} not supported!")
161
-
162
-
163
- def output_fn(output_df, accept_type):
164
- """Supports both CSV and JSON output formats."""
165
- if "text/csv" in accept_type:
166
- csv_output = output_df.fillna("N/A").to_csv(index=False) # CSV with N/A for missing values
167
- return csv_output, "text/csv"
168
- elif "application/json" in accept_type:
169
- return output_df.to_json(orient="records"), "application/json" # JSON array of records (NaNs -> null)
170
- else:
171
- raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
172
-
173
-
174
- def predict_fn(df, model):
175
- """Make predictions using MAPIE CQR and return the DataFrame with results."""
176
- model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
177
-
178
- # Load feature columns from the saved file
179
- with open(os.path.join(model_dir, "feature_columns.json")) as fp:
180
- model_features = json.load(fp)
181
-
182
- # Match features in a case-insensitive manner
183
- matched_df = match_features_case_insensitive(df, model_features)
184
-
185
- # Get CQR predictions - returns point prediction and intervals
186
- X_pred = matched_df[model_features]
187
- y_pred, y_pis = model.predict(X_pred)
188
-
189
- # Add predictions to dataframe with 95% intervals
190
- df["prediction"] = y_pred
191
- df["q_025"] = y_pis[:, 0, 0] # Lower bound (2.5th percentile)
192
- df["q_975"] = y_pis[:, 1, 0] # Upper bound (97.5th percentile)
193
-
194
- # Calculate std estimate from 95% interval
195
- interval_width_95 = df["q_975"] - df["q_025"]
196
- df["prediction_std"] = interval_width_95 / 3.92 # 95% CI = ±1.96σ, so width = 3.92σ
197
-
198
- # Calculate 50% intervals using normal approximation
199
- df["q_25"] = df["prediction"] - 0.674 * df["prediction_std"]
200
- df["q_75"] = df["prediction"] + 0.674 * df["prediction_std"]
201
-
202
- # Return the modified DataFrame
203
- return df