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.
- workbench/algorithms/dataframe/__init__.py +1 -2
- workbench/algorithms/dataframe/compound_dataset_overlap.py +321 -0
- workbench/algorithms/dataframe/feature_space_proximity.py +168 -75
- workbench/algorithms/dataframe/fingerprint_proximity.py +422 -86
- workbench/algorithms/dataframe/projection_2d.py +44 -21
- workbench/algorithms/dataframe/proximity.py +259 -305
- workbench/algorithms/graph/light/proximity_graph.py +14 -12
- workbench/algorithms/models/cleanlab_model.py +382 -0
- workbench/algorithms/models/noise_model.py +388 -0
- workbench/algorithms/sql/outliers.py +3 -3
- workbench/api/__init__.py +5 -1
- workbench/api/compound.py +1 -1
- workbench/api/df_store.py +17 -108
- workbench/api/endpoint.py +18 -5
- workbench/api/feature_set.py +121 -15
- workbench/api/meta.py +5 -2
- workbench/api/meta_model.py +289 -0
- workbench/api/model.py +55 -21
- workbench/api/monitor.py +1 -16
- workbench/api/parameter_store.py +3 -52
- workbench/cached/cached_model.py +4 -4
- workbench/core/artifacts/__init__.py +11 -2
- workbench/core/artifacts/artifact.py +16 -8
- workbench/core/artifacts/data_capture_core.py +355 -0
- workbench/core/artifacts/df_store_core.py +114 -0
- workbench/core/artifacts/endpoint_core.py +382 -253
- workbench/core/artifacts/feature_set_core.py +249 -45
- workbench/core/artifacts/model_core.py +135 -80
- workbench/core/artifacts/monitor_core.py +33 -248
- workbench/core/artifacts/parameter_store_core.py +98 -0
- workbench/core/cloud_platform/aws/aws_account_clamp.py +50 -1
- workbench/core/cloud_platform/aws/aws_meta.py +12 -5
- workbench/core/cloud_platform/aws/aws_session.py +4 -4
- workbench/core/pipelines/pipeline_executor.py +1 -1
- workbench/core/transforms/data_to_features/light/molecular_descriptors.py +4 -4
- workbench/core/transforms/features_to_model/features_to_model.py +62 -40
- workbench/core/transforms/model_to_endpoint/model_to_endpoint.py +76 -15
- workbench/core/transforms/pandas_transforms/pandas_to_features.py +38 -2
- workbench/core/views/training_view.py +113 -42
- workbench/core/views/view.py +53 -3
- workbench/core/views/view_utils.py +4 -4
- workbench/model_script_utils/model_script_utils.py +339 -0
- workbench/model_script_utils/pytorch_utils.py +405 -0
- workbench/model_script_utils/uq_harness.py +278 -0
- workbench/model_scripts/chemprop/chemprop.template +649 -0
- workbench/model_scripts/chemprop/generated_model_script.py +649 -0
- workbench/model_scripts/chemprop/model_script_utils.py +339 -0
- workbench/model_scripts/chemprop/requirements.txt +3 -0
- workbench/model_scripts/custom_models/chem_info/fingerprints.py +175 -0
- workbench/model_scripts/custom_models/chem_info/mol_descriptors.py +483 -0
- workbench/model_scripts/custom_models/chem_info/mol_standardize.py +450 -0
- workbench/model_scripts/custom_models/chem_info/molecular_descriptors.py +7 -9
- workbench/model_scripts/custom_models/chem_info/morgan_fingerprints.py +1 -1
- workbench/model_scripts/custom_models/proximity/feature_space_proximity.py +194 -0
- workbench/model_scripts/custom_models/proximity/feature_space_proximity.template +8 -10
- workbench/model_scripts/custom_models/uq_models/bayesian_ridge.template +7 -8
- workbench/model_scripts/custom_models/uq_models/ensemble_xgb.template +20 -21
- workbench/model_scripts/custom_models/uq_models/feature_space_proximity.py +194 -0
- workbench/model_scripts/custom_models/uq_models/gaussian_process.template +5 -11
- workbench/model_scripts/custom_models/uq_models/ngboost.template +30 -18
- workbench/model_scripts/custom_models/uq_models/requirements.txt +1 -3
- workbench/model_scripts/ensemble_xgb/ensemble_xgb.template +15 -17
- workbench/model_scripts/meta_model/generated_model_script.py +209 -0
- workbench/model_scripts/meta_model/meta_model.template +209 -0
- workbench/model_scripts/pytorch_model/generated_model_script.py +444 -500
- workbench/model_scripts/pytorch_model/model_script_utils.py +339 -0
- workbench/model_scripts/pytorch_model/pytorch.template +440 -496
- workbench/model_scripts/pytorch_model/pytorch_utils.py +405 -0
- workbench/model_scripts/pytorch_model/requirements.txt +1 -1
- workbench/model_scripts/pytorch_model/uq_harness.py +278 -0
- workbench/model_scripts/scikit_learn/generated_model_script.py +7 -12
- workbench/model_scripts/scikit_learn/scikit_learn.template +4 -9
- workbench/model_scripts/script_generation.py +20 -11
- workbench/model_scripts/uq_models/generated_model_script.py +248 -0
- workbench/model_scripts/xgb_model/generated_model_script.py +372 -404
- workbench/model_scripts/xgb_model/model_script_utils.py +339 -0
- workbench/model_scripts/xgb_model/uq_harness.py +278 -0
- workbench/model_scripts/xgb_model/xgb_model.template +369 -401
- workbench/repl/workbench_shell.py +28 -19
- workbench/resources/open_source_api.key +1 -1
- workbench/scripts/endpoint_test.py +162 -0
- workbench/scripts/lambda_test.py +73 -0
- workbench/scripts/meta_model_sim.py +35 -0
- workbench/scripts/ml_pipeline_batch.py +137 -0
- workbench/scripts/ml_pipeline_sqs.py +186 -0
- workbench/scripts/monitor_cloud_watch.py +20 -100
- workbench/scripts/training_test.py +85 -0
- workbench/utils/aws_utils.py +4 -3
- workbench/utils/chem_utils/__init__.py +0 -0
- workbench/utils/chem_utils/fingerprints.py +175 -0
- workbench/utils/chem_utils/misc.py +194 -0
- workbench/utils/chem_utils/mol_descriptors.py +483 -0
- workbench/utils/chem_utils/mol_standardize.py +450 -0
- workbench/utils/chem_utils/mol_tagging.py +348 -0
- workbench/utils/chem_utils/projections.py +219 -0
- workbench/utils/chem_utils/salts.py +256 -0
- workbench/utils/chem_utils/sdf.py +292 -0
- workbench/utils/chem_utils/toxicity.py +250 -0
- workbench/utils/chem_utils/vis.py +253 -0
- workbench/utils/chemprop_utils.py +141 -0
- workbench/utils/cloudwatch_handler.py +1 -1
- workbench/utils/cloudwatch_utils.py +137 -0
- workbench/utils/config_manager.py +3 -7
- workbench/utils/endpoint_utils.py +5 -7
- workbench/utils/license_manager.py +2 -6
- workbench/utils/meta_model_simulator.py +499 -0
- workbench/utils/metrics_utils.py +256 -0
- workbench/utils/model_utils.py +278 -79
- workbench/utils/monitor_utils.py +44 -62
- workbench/utils/pandas_utils.py +3 -3
- workbench/utils/pytorch_utils.py +87 -0
- workbench/utils/shap_utils.py +11 -57
- workbench/utils/workbench_logging.py +0 -3
- workbench/utils/workbench_sqs.py +1 -1
- workbench/utils/xgboost_local_crossfold.py +267 -0
- workbench/utils/xgboost_model_utils.py +127 -219
- workbench/web_interface/components/model_plot.py +14 -2
- workbench/web_interface/components/plugin_unit_test.py +5 -2
- workbench/web_interface/components/plugins/dashboard_status.py +3 -1
- workbench/web_interface/components/plugins/generated_compounds.py +1 -1
- workbench/web_interface/components/plugins/model_details.py +38 -74
- workbench/web_interface/components/plugins/scatter_plot.py +6 -10
- {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/METADATA +31 -9
- {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/RECORD +128 -96
- workbench-0.8.220.dist-info/entry_points.txt +11 -0
- {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/licenses/LICENSE +1 -1
- workbench/core/cloud_platform/aws/aws_df_store.py +0 -404
- workbench/core/cloud_platform/aws/aws_parameter_store.py +0 -280
- workbench/model_scripts/custom_models/chem_info/local_utils.py +0 -769
- workbench/model_scripts/custom_models/chem_info/tautomerize.py +0 -83
- workbench/model_scripts/custom_models/meta_endpoints/example.py +0 -53
- workbench/model_scripts/custom_models/proximity/generated_model_script.py +0 -138
- workbench/model_scripts/custom_models/proximity/proximity.py +0 -384
- workbench/model_scripts/custom_models/uq_models/generated_model_script.py +0 -393
- workbench/model_scripts/custom_models/uq_models/mapie_xgb.template +0 -203
- workbench/model_scripts/custom_models/uq_models/meta_uq.template +0 -273
- workbench/model_scripts/custom_models/uq_models/proximity.py +0 -384
- workbench/model_scripts/ensemble_xgb/generated_model_script.py +0 -279
- workbench/model_scripts/quant_regression/quant_regression.template +0 -279
- workbench/model_scripts/quant_regression/requirements.txt +0 -1
- workbench/utils/chem_utils.py +0 -1556
- workbench/utils/execution_environment.py +0 -211
- workbench/utils/fast_inference.py +0 -167
- workbench/utils/resource_utils.py +0 -39
- workbench-0.8.162.dist-info/entry_points.txt +0 -5
- {workbench-0.8.162.dist-info → workbench-0.8.220.dist-info}/WHEEL +0 -0
- {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
|