workbench 0.8.189__py3-none-any.whl → 0.8.190__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.

@@ -1,468 +0,0 @@
1
- # Imports for XGB Model
2
- import xgboost as xgb
3
- import awswrangler as wr
4
- import numpy as np
5
-
6
- # Model Performance Scores
7
- from sklearn.metrics import (
8
- mean_absolute_error,
9
- r2_score,
10
- root_mean_squared_error,
11
- precision_recall_fscore_support,
12
- confusion_matrix,
13
- )
14
-
15
- # Classification Encoder
16
- from sklearn.preprocessing import LabelEncoder
17
-
18
- # Scikit Learn Imports
19
- from sklearn.model_selection import train_test_split
20
-
21
- from io import StringIO
22
- import json
23
- import argparse
24
- import joblib
25
- import os
26
- import pandas as pd
27
- from typing import List, Tuple
28
-
29
- # Template Parameters
30
- TEMPLATE_PARAMS = {
31
- "model_type": "regressor",
32
- "target": "class_number_of_rings",
33
- "features": ['length', 'diameter', 'height', 'whole_weight', 'shucked_weight', 'viscera_weight', 'shell_weight', 'sex'],
34
- "compressed_features": [],
35
- "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/abalone-regression/training",
36
- "train_all_data": False,
37
- "hyperparameters": {},
38
- }
39
-
40
-
41
- # Function to check if dataframe is empty
42
- def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
43
- """
44
- Check if the provided dataframe is empty and raise an exception if it is.
45
-
46
- Args:
47
- df (pd.DataFrame): DataFrame to check
48
- df_name (str): Name of the DataFrame
49
- """
50
- if df.empty:
51
- msg = f"*** The training data {df_name} has 0 rows! ***STOPPING***"
52
- print(msg)
53
- raise ValueError(msg)
54
-
55
-
56
- def expand_proba_column(df: pd.DataFrame, class_labels: List[str]) -> pd.DataFrame:
57
- """
58
- Expands a column in a DataFrame containing a list of probabilities into separate columns.
59
-
60
- Args:
61
- df (pd.DataFrame): DataFrame containing a "pred_proba" column
62
- class_labels (List[str]): List of class labels
63
-
64
- Returns:
65
- pd.DataFrame: DataFrame with the "pred_proba" expanded into separate columns
66
- """
67
-
68
- # Sanity check
69
- proba_column = "pred_proba"
70
- if proba_column not in df.columns:
71
- raise ValueError('DataFrame does not contain a "pred_proba" column')
72
-
73
- # Construct new column names with '_proba' suffix
74
- proba_splits = [f"{label}_proba" for label in class_labels]
75
-
76
- # Expand the proba_column into separate columns for each probability
77
- proba_df = pd.DataFrame(df[proba_column].tolist(), columns=proba_splits)
78
-
79
- # Drop any proba columns and reset the index in prep for the concat
80
- df = df.drop(columns=[proba_column] + proba_splits, errors="ignore")
81
- df = df.reset_index(drop=True)
82
-
83
- # Concatenate the new columns with the original DataFrame
84
- df = pd.concat([df, proba_df], axis=1)
85
- print(df)
86
- return df
87
-
88
-
89
- def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
90
- """
91
- Matches and renames DataFrame columns to match model feature names (case-insensitive).
92
- Prioritizes exact matches, then case-insensitive matches.
93
-
94
- Raises ValueError if any model features cannot be matched.
95
- """
96
- df_columns_lower = {col.lower(): col for col in df.columns}
97
- rename_dict = {}
98
- missing = []
99
- for feature in model_features:
100
- if feature in df.columns:
101
- continue # Exact match
102
- elif feature.lower() in df_columns_lower:
103
- rename_dict[df_columns_lower[feature.lower()]] = feature
104
- else:
105
- missing.append(feature)
106
-
107
- if missing:
108
- raise ValueError(f"Features not found: {missing}")
109
-
110
- # Rename the DataFrame columns to match the model features
111
- return df.rename(columns=rename_dict)
112
-
113
-
114
- def convert_categorical_types(df: pd.DataFrame, features: list, category_mappings={}) -> tuple:
115
- """
116
- Converts appropriate columns to categorical type with consistent mappings.
117
-
118
- Args:
119
- df (pd.DataFrame): The DataFrame to process.
120
- features (list): List of feature names to consider for conversion.
121
- category_mappings (dict, optional): Existing category mappings. If empty dict, we're in
122
- training mode. If populated, we're in inference mode.
123
-
124
- Returns:
125
- tuple: (processed DataFrame, category mappings dictionary)
126
- """
127
- # Training mode
128
- if category_mappings == {}:
129
- for col in df.select_dtypes(include=["object", "string"]):
130
- if col in features and df[col].nunique() < 20:
131
- print(f"Training mode: Converting {col} to category")
132
- df[col] = df[col].astype("category")
133
- category_mappings[col] = df[col].cat.categories.tolist() # Store category mappings
134
-
135
- # Inference mode
136
- else:
137
- for col, categories in category_mappings.items():
138
- if col in df.columns:
139
- print(f"Inference mode: Applying categorical mapping for {col}")
140
- df[col] = pd.Categorical(df[col], categories=categories) # Apply consistent categorical mapping
141
-
142
- return df, category_mappings
143
-
144
-
145
- def decompress_features(
146
- df: pd.DataFrame, features: List[str], compressed_features: List[str]
147
- ) -> Tuple[pd.DataFrame, List[str]]:
148
- """Prepare features for the model by decompressing bitstring features
149
-
150
- Args:
151
- df (pd.DataFrame): The features DataFrame
152
- features (List[str]): Full list of feature names
153
- compressed_features (List[str]): List of feature names to decompress (bitstrings)
154
-
155
- Returns:
156
- pd.DataFrame: DataFrame with the decompressed features
157
- List[str]: Updated list of feature names after decompression
158
-
159
- Raises:
160
- ValueError: If any missing values are found in the specified features
161
- """
162
-
163
- # Check for any missing values in the required features
164
- missing_counts = df[features].isna().sum()
165
- if missing_counts.any():
166
- missing_features = missing_counts[missing_counts > 0]
167
- print(
168
- f"WARNING: Found missing values in features: {missing_features.to_dict()}. "
169
- "WARNING: You might want to remove/replace all NaN values before processing."
170
- )
171
-
172
- # Decompress the specified compressed features
173
- decompressed_features = features.copy()
174
- for feature in compressed_features:
175
- if (feature not in df.columns) or (feature not in features):
176
- print(f"Feature '{feature}' not in the features list, skipping decompression.")
177
- continue
178
-
179
- # Remove the feature from the list of features to avoid duplication
180
- decompressed_features.remove(feature)
181
-
182
- # Handle all compressed features as bitstrings
183
- bit_matrix = np.array([list(bitstring) for bitstring in df[feature]], dtype=np.uint8)
184
- prefix = feature[:3]
185
-
186
- # Create all new columns at once - avoids fragmentation
187
- new_col_names = [f"{prefix}_{i}" for i in range(bit_matrix.shape[1])]
188
- new_df = pd.DataFrame(bit_matrix, columns=new_col_names, index=df.index)
189
-
190
- # Add to features list
191
- decompressed_features.extend(new_col_names)
192
-
193
- # Drop original column and concatenate new ones
194
- df = df.drop(columns=[feature])
195
- df = pd.concat([df, new_df], axis=1)
196
-
197
- return df, decompressed_features
198
-
199
-
200
- if __name__ == "__main__":
201
- """The main function is for training the XGBoost model"""
202
-
203
- # Harness Template Parameters
204
- target = TEMPLATE_PARAMS["target"]
205
- features = TEMPLATE_PARAMS["features"]
206
- orig_features = features.copy()
207
- compressed_features = TEMPLATE_PARAMS["compressed_features"]
208
- model_type = TEMPLATE_PARAMS["model_type"]
209
- model_metrics_s3_path = TEMPLATE_PARAMS["model_metrics_s3_path"]
210
- train_all_data = TEMPLATE_PARAMS["train_all_data"]
211
- hyperparameters = TEMPLATE_PARAMS["hyperparameters"]
212
- validation_split = 0.2
213
-
214
- # Script arguments for input/output directories
215
- parser = argparse.ArgumentParser()
216
- parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
217
- parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
218
- parser.add_argument(
219
- "--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data")
220
- )
221
- args = parser.parse_args()
222
-
223
- # Read the training data into DataFrames
224
- training_files = [os.path.join(args.train, file) for file in os.listdir(args.train) if file.endswith(".csv")]
225
- print(f"Training Files: {training_files}")
226
-
227
- # Combine files and read them all into a single pandas dataframe
228
- all_df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
229
-
230
- # Check if the dataframe is empty
231
- check_dataframe(all_df, "training_df")
232
-
233
- # Features/Target output
234
- print(f"Target: {target}")
235
- print(f"Features: {str(features)}")
236
-
237
- # Convert any features that might be categorical to 'category' type
238
- all_df, category_mappings = convert_categorical_types(all_df, features)
239
-
240
- # If we have compressed features, decompress them
241
- if compressed_features:
242
- print(f"Decompressing features {compressed_features}...")
243
- all_df, features = decompress_features(all_df, features, compressed_features)
244
-
245
- # Do we want to train on all the data?
246
- if train_all_data:
247
- print("Training on ALL of the data")
248
- df_train = all_df.copy()
249
- df_val = all_df.copy()
250
-
251
- # Does the dataframe have a training column?
252
- elif "training" in all_df.columns:
253
- print("Found training column, splitting data based on training column")
254
- df_train = all_df[all_df["training"]]
255
- df_val = all_df[~all_df["training"]]
256
- else:
257
- # Just do a random training Split
258
- print("WARNING: No training column found, splitting data with random state=42")
259
- df_train, df_val = train_test_split(all_df, test_size=validation_split, random_state=42)
260
- print(f"FIT/TRAIN: {df_train.shape}")
261
- print(f"VALIDATION: {df_val.shape}")
262
-
263
- # Use any hyperparameters to set up both the trainer and model configurations
264
- print(f"Hyperparameters: {hyperparameters}")
265
-
266
- # Now spin up our XGB Model
267
- if model_type == "classifier":
268
- xgb_model = xgb.XGBClassifier(enable_categorical=True, **hyperparameters)
269
-
270
- # Encode the target column
271
- label_encoder = LabelEncoder()
272
- df_train[target] = label_encoder.fit_transform(df_train[target])
273
- df_val[target] = label_encoder.transform(df_val[target])
274
-
275
- else:
276
- xgb_model = xgb.XGBRegressor(enable_categorical=True, **hyperparameters)
277
- label_encoder = None # We don't need this for regression
278
-
279
- # Grab our Features, Target and Train the Model
280
- y_train = df_train[target]
281
- X_train = df_train[features]
282
- xgb_model.fit(X_train, y_train)
283
-
284
- # Make Predictions on the Validation Set
285
- print(f"Making Predictions on Validation Set...")
286
- y_validate = df_val[target]
287
- X_validate = df_val[features]
288
- preds = xgb_model.predict(X_validate)
289
- if model_type == "classifier":
290
- # Also get the probabilities for each class
291
- print("Processing Probabilities...")
292
- probs = xgb_model.predict_proba(X_validate)
293
- df_val["pred_proba"] = [p.tolist() for p in probs]
294
-
295
- # Expand the pred_proba column into separate columns for each class
296
- print(df_val.columns)
297
- df_val = expand_proba_column(df_val, label_encoder.classes_)
298
- print(df_val.columns)
299
-
300
- # Decode the target and prediction labels
301
- y_validate = label_encoder.inverse_transform(y_validate)
302
- preds = label_encoder.inverse_transform(preds)
303
-
304
- # Save predictions to S3 (just the target, prediction, and '_proba' columns)
305
- df_val["prediction"] = preds
306
- output_columns = [target, "prediction"]
307
- output_columns += [col for col in df_val.columns if col.endswith("_proba")]
308
- wr.s3.to_csv(
309
- df_val[output_columns],
310
- path=f"{model_metrics_s3_path}/validation_predictions.csv",
311
- index=False,
312
- )
313
-
314
- # Report Performance Metrics
315
- if model_type == "classifier":
316
- # Get the label names and their integer mapping
317
- label_names = label_encoder.classes_
318
-
319
- # Calculate various model performance metrics
320
- scores = precision_recall_fscore_support(y_validate, preds, average=None, labels=label_names)
321
-
322
- # Put the scores into a dataframe
323
- score_df = pd.DataFrame(
324
- {
325
- target: label_names,
326
- "precision": scores[0],
327
- "recall": scores[1],
328
- "fscore": scores[2],
329
- "support": scores[3],
330
- }
331
- )
332
-
333
- # We need to get creative with the Classification Metrics
334
- metrics = ["precision", "recall", "fscore", "support"]
335
- for t in label_names:
336
- for m in metrics:
337
- value = score_df.loc[score_df[target] == t, m].iloc[0]
338
- print(f"Metrics:{t}:{m} {value}")
339
-
340
- # Compute and output the confusion matrix
341
- conf_mtx = confusion_matrix(y_validate, preds, labels=label_names)
342
- for i, row_name in enumerate(label_names):
343
- for j, col_name in enumerate(label_names):
344
- value = conf_mtx[i, j]
345
- print(f"ConfusionMatrix:{row_name}:{col_name} {value}")
346
-
347
- else:
348
- # Calculate various model performance metrics (regression)
349
- rmse = root_mean_squared_error(y_validate, preds)
350
- mae = mean_absolute_error(y_validate, preds)
351
- r2 = r2_score(y_validate, preds)
352
- print(f"RMSE: {rmse:.3f}")
353
- print(f"MAE: {mae:.3f}")
354
- print(f"R2: {r2:.3f}")
355
- print(f"NumRows: {len(df_val)}")
356
-
357
- # Now save the model to the standard place/name
358
- joblib.dump(xgb_model, os.path.join(args.model_dir, "xgb_model.joblib"))
359
-
360
- # Save the label encoder if we have one
361
- if label_encoder:
362
- joblib.dump(label_encoder, os.path.join(args.model_dir, "label_encoder.joblib"))
363
-
364
- # Save the features (this will validate input during predictions)
365
- with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
366
- json.dump(orig_features, fp) # We save the original features, not the decompressed ones
367
-
368
- # Save the category mappings
369
- with open(os.path.join(args.model_dir, "category_mappings.json"), "w") as fp:
370
- json.dump(category_mappings, fp)
371
-
372
-
373
- def model_fn(model_dir):
374
- """Deserialize and return fitted XGBoost model"""
375
- model_path = os.path.join(model_dir, "xgb_model.joblib")
376
- model = joblib.load(model_path)
377
- return model
378
-
379
-
380
- def input_fn(input_data, content_type):
381
- """Parse input data and return a DataFrame."""
382
- if not input_data:
383
- raise ValueError("Empty input data is not supported!")
384
-
385
- # Decode bytes to string if necessary
386
- if isinstance(input_data, bytes):
387
- input_data = input_data.decode("utf-8")
388
-
389
- if "text/csv" in content_type:
390
- return pd.read_csv(StringIO(input_data))
391
- elif "application/json" in content_type:
392
- return pd.DataFrame(json.loads(input_data)) # Assumes JSON array of records
393
- else:
394
- raise ValueError(f"{content_type} not supported!")
395
-
396
-
397
- def output_fn(output_df, accept_type):
398
- """Supports both CSV and JSON output formats."""
399
- if "text/csv" in accept_type:
400
- csv_output = output_df.fillna("N/A").to_csv(index=False) # CSV with N/A for missing values
401
- return csv_output, "text/csv"
402
- elif "application/json" in accept_type:
403
- return output_df.to_json(orient="records"), "application/json" # JSON array of records (NaNs -> null)
404
- else:
405
- raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
406
-
407
-
408
- def predict_fn(df, model) -> pd.DataFrame:
409
- """Make Predictions with our XGB Model
410
-
411
- Args:
412
- df (pd.DataFrame): The input DataFrame
413
- model: The model use for predictions
414
-
415
- Returns:
416
- pd.DataFrame: The DataFrame with the predictions added
417
- """
418
- compressed_features = TEMPLATE_PARAMS["compressed_features"]
419
-
420
- # Grab our feature columns (from training)
421
- model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
422
- with open(os.path.join(model_dir, "feature_columns.json")) as fp:
423
- features = json.load(fp)
424
- print(f"Model Features: {features}")
425
-
426
- # Load the category mappings (from training)
427
- with open(os.path.join(model_dir, "category_mappings.json")) as fp:
428
- category_mappings = json.load(fp)
429
-
430
- # Load our Label Encoder if we have one
431
- label_encoder = None
432
- if os.path.exists(os.path.join(model_dir, "label_encoder.joblib")):
433
- label_encoder = joblib.load(os.path.join(model_dir, "label_encoder.joblib"))
434
-
435
- # We're going match features in a case-insensitive manner, accounting for all the permutations
436
- # - Model has a feature list that's any case ("Id", "taCos", "cOunT", "likes_tacos")
437
- # - Incoming data has columns that are mixed case ("ID", "Tacos", "Count", "Likes_Tacos")
438
- matched_df = match_features_case_insensitive(df, features)
439
-
440
- # Detect categorical types in the incoming DataFrame
441
- matched_df, _ = convert_categorical_types(matched_df, features, category_mappings)
442
-
443
- # If we have compressed features, decompress them
444
- if compressed_features:
445
- print("Decompressing features for prediction...")
446
- matched_df, features = decompress_features(matched_df, features, compressed_features)
447
-
448
- # Predict the features against our XGB Model
449
- X = matched_df[features]
450
- predictions = model.predict(X)
451
-
452
- # If we have a label encoder, decode the predictions
453
- if label_encoder:
454
- predictions = label_encoder.inverse_transform(predictions)
455
-
456
- # Set the predictions on the DataFrame
457
- df["prediction"] = predictions
458
-
459
- # Does our model have a 'predict_proba' method? If so we will call it and add the results to the DataFrame
460
- if getattr(model, "predict_proba", None):
461
- probs = model.predict_proba(matched_df[features])
462
- df["pred_proba"] = [p.tolist() for p in probs]
463
-
464
- # Expand the pred_proba column into separate columns for each class
465
- df = expand_proba_column(df, label_encoder.classes_)
466
-
467
- # All done, return the DataFrame with new columns for the predictions
468
- return df