workbench 0.8.198__py3-none-any.whl → 0.8.201__py3-none-any.whl

workbench/model_scripts/pytorch_model/pytorch.template

@@ -29,9 +29,7 @@ from io import StringIO
 import json
 import argparse
 import joblib
-import os
 import pandas as pd
-from typing import List, Tuple
 
 # Template Parameters
 TEMPLATE_PARAMS = {
@@ -45,7 +43,6 @@ TEMPLATE_PARAMS = {
 }
 
 
-# Function to check if dataframe is empty
 def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
     """
     Check if the provided dataframe is empty and raise an exception if it is.
@@ -60,19 +57,17 @@ def check_dataframe(df: pd.DataFrame, df_name: str) -> None:
         raise ValueError(msg)
 
 
-def expand_proba_column(df: pd.DataFrame, class_labels: List[str]) -> pd.DataFrame:
+def expand_proba_column(df: pd.DataFrame, class_labels: list[str]) -> pd.DataFrame:
     """
     Expands a column in a DataFrame containing a list of probabilities into separate columns.
 
     Args:
         df (pd.DataFrame): DataFrame containing a "pred_proba" column
-        class_labels (List[str]): List of class labels
+        class_labels (list[str]): List of class labels
 
     Returns:
         pd.DataFrame: DataFrame with the "pred_proba" expanded into separate columns
     """
-
-    # Sanity check
     proba_column = "pred_proba"
     if proba_column not in df.columns:
         raise ValueError('DataFrame does not contain a "pred_proba" column')
@@ -89,11 +84,10 @@ def expand_proba_column(df: pd.DataFrame, class_labels: List[str]) -> pd.DataFra
 
     # Concatenate the new columns with the original DataFrame
     df = pd.concat([df, proba_df], axis=1)
-    print(df)
     return df
 
 
-def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
+def match_features_case_insensitive(df: pd.DataFrame, model_features: list[str]) -> pd.DataFrame:
     """
     Matches and renames DataFrame columns to match model feature names (case-insensitive).
     Prioritizes exact matches, then case-insensitive matches.
@@ -118,55 +112,60 @@ def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> p
     return df.rename(columns=rename_dict)
 
 
-def convert_categorical_types(df: pd.DataFrame, features: list, category_mappings={}) -> tuple:
+def convert_categorical_types(
+    df: pd.DataFrame, features: list[str], category_mappings: dict[str, list[str]] | None = None
+) -> tuple[pd.DataFrame, dict[str, list[str]]]:
     """
     Converts appropriate columns to categorical type with consistent mappings.
 
     Args:
         df (pd.DataFrame): The DataFrame to process.
         features (list): List of feature names to consider for conversion.
-        category_mappings (dict, optional): Existing category mappings. If empty dict, we're in
-                                            training mode. If populated, we're in inference mode.
+        category_mappings (dict, optional): Existing category mappings. If None or empty,
+                                            we're in training mode. If populated, we're in
+                                            inference mode.
 
     Returns:
         tuple: (processed DataFrame, category mappings dictionary)
     """
+    if category_mappings is None:
+        category_mappings = {}
+
     # Training mode
-    if category_mappings == {}:
+    if not category_mappings:
         for col in df.select_dtypes(include=["object", "string"]):
             if col in features and df[col].nunique() < 20:
                 print(f"Training mode: Converting {col} to category")
                 df[col] = df[col].astype("category")
-                category_mappings[col] = df[col].cat.categories.tolist()  # Store category mappings
+                category_mappings[col] = df[col].cat.categories.tolist()
 
     # Inference mode
     else:
         for col, categories in category_mappings.items():
             if col in df.columns:
                 print(f"Inference mode: Applying categorical mapping for {col}")
-                df[col] = pd.Categorical(df[col], categories=categories)  # Apply consistent categorical mapping
+                df[col] = pd.Categorical(df[col], categories=categories)
 
     return df, category_mappings
 
 
 def decompress_features(
-    df: pd.DataFrame, features: List[str], compressed_features: List[str]
-) -> Tuple[pd.DataFrame, List[str]]:
+    df: pd.DataFrame, features: list[str], compressed_features: list[str]
+) -> tuple[pd.DataFrame, list[str]]:
     """Prepare features for the model
 
     Args:
         df (pd.DataFrame): The features DataFrame
-        features (List[str]): Full list of feature names
-        compressed_features (List[str]): List of feature names to decompress (bitstrings)
+        features (list[str]): Full list of feature names
+        compressed_features (list[str]): List of feature names to decompress (bitstrings)
 
     Returns:
         pd.DataFrame: DataFrame with the decompressed features
-        List[str]: Updated list of feature names after decompression
+        list[str]: Updated list of feature names after decompression
 
     Raises:
         ValueError: If any missing values are found in the specified features
     """
-
     # Check for any missing values in the required features
     missing_counts = df[features].isna().sum()
     if missing_counts.any():
@@ -176,10 +175,11 @@ def decompress_features(
             "WARNING: You might want to remove/replace all NaN values before processing."
         )
 
-    # Decompress the specified compressed features
-    decompressed_features = features
+    # Make a copy to avoid mutating the original list
+    decompressed_features = features.copy()
+
     for feature in compressed_features:
-        if (feature not in df.columns) or (feature not in features):
+        if (feature not in df.columns) or (feature not in decompressed_features):
             print(f"Feature '{feature}' not in the features list, skipping decompression.")
             continue
 
@@ -204,26 +204,39 @@ def decompress_features(
     return df, decompressed_features
 
 
-def model_fn(model_dir):
+def model_fn(model_dir: str) -> TabularModel:
+    """Load the PyTorch Tabular model from the specified directory.
 
+    Args:
+        model_dir: Directory containing the saved model
+
+    Returns:
+        Loaded TabularModel instance
+    """
     # Save current working directory
     original_cwd = os.getcwd()
     try:
         # Change to /tmp because Pytorch Tabular needs write access (creates a .pt_tmp directory)
         os.chdir("/tmp")
 
-        # Load the model
+        # Remove callbacks.sav if it exists - it's not needed for inference and causes
+        # GPU->CPU loading issues (joblib.load doesn't support map_location)
         model_path = os.path.join(model_dir, "tabular_model")
-        model = TabularModel.load_model(model_path)
+        callbacks_path = os.path.join(model_path, "callbacks.sav")
+        if os.path.exists(callbacks_path):
+            os.remove(callbacks_path)
+
+        # Load the model (map_location="cpu" ensures GPU-trained models work on CPU endpoints)
+        model = TabularModel.load_model(model_path, map_location="cpu")
 
-    # Restore the original working directory
     finally:
+        # Restore the original working directory
         os.chdir(original_cwd)
 
     return model
 
 
-def input_fn(input_data, content_type):
+def input_fn(input_data, content_type: str) -> pd.DataFrame:
     """Parse input data and return a DataFrame."""
     if not input_data:
         raise ValueError("Empty input data is not supported!")
@@ -240,18 +253,18 @@ def input_fn(input_data, content_type):
         raise ValueError(f"{content_type} not supported!")
 
 
-def output_fn(output_df, accept_type):
+def output_fn(output_df: pd.DataFrame, accept_type: str) -> tuple[str, str]:
     """Supports both CSV and JSON output formats."""
     if "text/csv" in accept_type:
-        csv_output = output_df.fillna("N/A").to_csv(index=False)  # CSV with N/A for missing values
+        csv_output = output_df.fillna("N/A").to_csv(index=False)
         return csv_output, "text/csv"
     elif "application/json" in accept_type:
-        return output_df.to_json(orient="records"), "application/json"  # JSON array of records (NaNs -> null)
+        return output_df.to_json(orient="records"), "application/json"
     else:
         raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
 
 
-def predict_fn(df, model) -> pd.DataFrame:
+def predict_fn(df: pd.DataFrame, model: TabularModel) -> pd.DataFrame:
     """Make Predictions with our PyTorch Tabular Model
 
     Args:
@@ -275,12 +288,11 @@ def predict_fn(df, model) -> pd.DataFrame:
 
     # Load our Label Encoder if we have one
     label_encoder = None
-    if os.path.exists(os.path.join(model_dir, "label_encoder.joblib")):
-        label_encoder = joblib.load(os.path.join(model_dir, "label_encoder.joblib"))
+    label_encoder_path = os.path.join(model_dir, "label_encoder.joblib")
+    if os.path.exists(label_encoder_path):
+        label_encoder = joblib.load(label_encoder_path)
 
-    # We're going match features in a case-insensitive manner, accounting for all the permutations
-    # - Model has a feature list that's any case ("Id", "taCos", "cOunT", "likes_tacos")
-    # - Incoming data has columns that are mixed case ("ID", "Tacos", "Count", "Likes_Tacos")
+    # Match features in a case-insensitive manner
     matched_df = match_features_case_insensitive(df, features)
 
     # Detect categorical types in the incoming DataFrame
@@ -291,12 +303,25 @@ def predict_fn(df, model) -> pd.DataFrame:
         print("Decompressing features for prediction...")
         matched_df, features = decompress_features(matched_df, features, compressed_features)
 
+    # Track rows with missing features
+    missing_mask = matched_df[features].isna().any(axis=1)
+    if missing_mask.any():
+        print(f"Warning: {missing_mask.sum()} rows have missing features, will return NaN predictions")
+
+    # Initialize prediction column with NaN
+    df["prediction"] = np.nan
+
+    # Only predict on complete rows
+    complete_df = matched_df[~missing_mask]
+    if len(complete_df) == 0:
+        print("Warning: No complete rows to predict on")
+        return df
+
     # Make predictions using the TabularModel
-    result = model.predict(matched_df[features])
+    result = model.predict(complete_df[features])
 
     # pytorch-tabular returns predictions using f"{target}_prediction" column
-    # and classification probabilities in columns ending with "_probability"
-    target = TEMPLATE_PARAMS["target_column"]
+    target = TEMPLATE_PARAMS["target"]
     prediction_column = f"{target}_prediction"
     if prediction_column in result.columns:
         predictions = result[prediction_column].values
@@ -307,20 +332,23 @@ def predict_fn(df, model) -> pd.DataFrame:
     if label_encoder:
         predictions = label_encoder.inverse_transform(predictions.astype(int))
 
-    # Set the predictions on the DataFrame
-    df["prediction"] = predictions
+    # Set predictions only for complete rows
+    df.loc[~missing_mask, "prediction"] = predictions
 
     # For classification, get probabilities
     if label_encoder is not None:
         prob_cols = [col for col in result.columns if col.endswith("_probability")]
         if prob_cols:
             probs = result[prob_cols].values
-            df["pred_proba"] = [p.tolist() for p in probs]
+
+            # Build full proba Series with None for missing rows
+            all_proba = pd.Series([None] * len(df), index=df.index, dtype=object)
+            all_proba.loc[~missing_mask] = [p.tolist() for p in probs]
+            df["pred_proba"] = all_proba
 
             # Expand the pred_proba column into separate columns for each class
             df = expand_proba_column(df, label_encoder.classes_)
 
-    # All done, return the DataFrame with new columns for the predictions
     return df
 
 
@@ -354,9 +382,21 @@ if __name__ == "__main__":
     # Combine files and read them all into a single pandas dataframe
     all_df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
 
+    # Print out some info about the dataframe
+    print(f"All Data Shape: {all_df.shape}")
+    print(f"Feature dtypes:\n{all_df[features].dtypes.value_counts()}")
+    print(f"Int64 columns: {all_df[features].select_dtypes(include=['int64']).columns.tolist()}")
+
     # Check if the dataframe is empty
     check_dataframe(all_df, "training_df")
 
+    # Drop any rows with missing feature values
+    initial_row_count = all_df.shape[0]
+    all_df = all_df.dropna(subset=features)
+    dropped_rows = initial_row_count - all_df.shape[0]
+    if dropped_rows > 0:
+        print(f"Dropped {dropped_rows} rows due to missing feature values.")
+
     # Features/Target output
     print(f"Target: {target}")
     print(f"Features: {str(features)}")
@@ -364,11 +404,25 @@ if __name__ == "__main__":
     # Convert any features that might be categorical to 'category' type
     all_df, category_mappings = convert_categorical_types(all_df, features)
 
+    # Print out some info about the dataframe
+    print(f"All Data Shape: {all_df.shape}")
+    print(f"Feature dtypes:\n{all_df[features].dtypes.value_counts()}")
+    print(f"Int64 columns: {all_df[features].select_dtypes(include=['int64']).columns.tolist()}")
+
     # If we have compressed features, decompress them
     if compressed_features:
         print(f"Decompressing features {compressed_features}...")
         all_df, features = decompress_features(all_df, features, compressed_features)
 
+    # Determine categorical and continuous columns
+    categorical_cols = [col for col in features if all_df[col].dtype.name == "category"]
+    continuous_cols = [col for col in features if col not in categorical_cols]
+    print(f"Categorical columns: {categorical_cols}")
+    print(f"Continuous columns: {continuous_cols}")
+
+    # Cast continuous columns to float
+    all_df[continuous_cols] = all_df[continuous_cols].astype("float64")
+
     # Do we want to train on all the data?
     if train_all_data:
         print("Training on ALL of the data")
@@ -378,8 +432,8 @@ if __name__ == "__main__":
     # Does the dataframe have a training column?
     elif "training" in all_df.columns:
         print("Found training column, splitting data based on training column")
-        df_train = all_df[all_df["training"]]
-        df_val = all_df[~all_df["training"]]
+        df_train = all_df[all_df["training"]].copy()
+        df_val = all_df[~all_df["training"]].copy()
     else:
         # Just do a random training Split
         print("WARNING: No training column found, splitting data with random state=42")
@@ -387,13 +441,6 @@ if __name__ == "__main__":
     print(f"FIT/TRAIN: {df_train.shape}")
     print(f"VALIDATION: {df_val.shape}")
 
-    # Determine categorical and continuous columns
-    categorical_cols = [col for col in features if df_train[col].dtype.name == "category"]
-    continuous_cols = [col for col in features if col not in categorical_cols]
-
-    print(f"Categorical columns: {categorical_cols}")
-    print(f"Continuous columns: {continuous_cols}")
-
     # Set up PyTorch Tabular configuration
     data_config = DataConfig(
         target=[target],
@@ -417,11 +464,12 @@ if __name__ == "__main__":
 
     # Set up PyTorch Tabular configuration with defaults
     trainer_defaults = {
-        "auto_lr_find": True,
-        "batch_size": min(1024, max(32, len(df_train) // 4)),
+        "auto_lr_find": False,
+        "batch_size": min(128, max(32, len(df_train) // 16)),
         "max_epochs": 100,
+        "min_epochs": 10,
         "early_stopping": "valid_loss",
-        "early_stopping_patience": 15,
+        "early_stopping_patience": 10,
         "checkpoints": "valid_loss",
         "accelerator": "auto",
         "progress_bar": "none",
@@ -430,7 +478,6 @@ if __name__ == "__main__":
 
     # 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}
-    # Print overwrites
     for key, value in training_overrides.items():
         print(f"TRAINING CONFIG Override: {key}: {trainer_defaults[key]} → {value}")
     trainer_params = {**trainer_defaults, **training_overrides}
@@ -438,23 +485,20 @@ if __name__ == "__main__":
 
     # Model config defaults
     model_defaults = {
-        "layers": "1024-512-512",
-        "activation": "ReLU",
+        "layers": "256-128-64",
+        "activation": "LeakyReLU",
         "learning_rate": 1e-3,
-        "dropout": 0.1,
+        "dropout": 0.3,
         "use_batch_norm": True,
         "initialization": "kaiming",
     }
     # 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}
-    # Print overwrites
     for key, value in model_overrides.items():
         print(f"MODEL CONFIG Override: {key}: {model_defaults[key]} → {value}")
     model_params = {**model_defaults, **model_overrides}
 
     # Use CategoryEmbedding model configuration for general-purpose tabular modeling.
-    # Works effectively for both regression and classification as the foundational
-    # architecture in PyTorch Tabular
     model_config = CategoryEmbeddingModelConfig(task=task, **model_params)
     optimizer_config = OptimizerConfig()
 
@@ -474,36 +518,34 @@ if __name__ == "__main__":
     result = tabular_model.predict(df_val, include_input_features=False)
 
     # pytorch-tabular returns predictions using f"{target}_prediction" column
-    # and classification probabilities in columns ending with "_probability"
-    if model_type == "classifier":
-        preds = result[f"{target}_prediction"].values
-    else:
-        # Regression: use the target column name
-        preds = result[f"{target}_prediction"].values
+    preds = result[f"{target}_prediction"].values
 
     if model_type == "classifier":
         # Get probabilities for classification
         print("Processing Probabilities...")
-        prob_cols = [col for col in result.columns if col.endswith("_probability")]
+        prob_cols = sorted([col for col in result.columns if col.endswith("_probability")])
         if prob_cols:
             probs = result[prob_cols].values
+            df_val = df_val.copy()  # Avoid SettingWithCopyWarning
             df_val["pred_proba"] = [p.tolist() for p in probs]
 
             # Expand the pred_proba column into separate columns for each class
-            print(df_val.columns)
+            print(df_val.columns.tolist())
             df_val = expand_proba_column(df_val, label_encoder.classes_)
-            print(df_val.columns)
+            print(df_val.columns.tolist())
 
         # Decode the target and prediction labels
         y_validate = label_encoder.inverse_transform(df_val[target])
-        preds = label_encoder.inverse_transform(preds.astype(int))
+        preds_decoded = label_encoder.inverse_transform(preds.astype(int))
     else:
         y_validate = df_val[target].values
+        preds_decoded = preds
 
-    # Save predictions to S3 (just the target, prediction, and '_probability' columns)
-    df_val["prediction"] = preds
+    # Save predictions to S3
+    df_val = df_val.copy()
+    df_val["prediction"] = preds_decoded
     output_columns = [target, "prediction"]
-    output_columns += [col for col in df_val.columns if col.endswith("_probability")]
+    output_columns += [col for col in df_val.columns if col.endswith("_proba")]
     wr.s3.to_csv(
         df_val[output_columns],
         path=f"{model_metrics_s3_path}/validation_predictions.csv",
@@ -516,7 +558,7 @@ if __name__ == "__main__":
         label_names = label_encoder.classes_
 
         # Calculate various model performance metrics
-        scores = precision_recall_fscore_support(y_validate, preds, average=None, labels=label_names)
+        scores = precision_recall_fscore_support(y_validate, preds_decoded, average=None, labels=label_names)
 
         # Put the scores into a dataframe
         score_df = pd.DataFrame(
@@ -524,20 +566,20 @@ if __name__ == "__main__":
                 target: label_names,
                 "precision": scores[0],
                 "recall": scores[1],
-                "fscore": scores[2],
+                "f1": scores[2],
                 "support": scores[3],
             }
         )
 
-        # We need to get creative with the Classification Metrics
-        metrics = ["precision", "recall", "fscore", "support"]
+        # Output metrics per class
+        metrics = ["precision", "recall", "f1", "support"]
         for t in label_names:
             for m in metrics:
                 value = score_df.loc[score_df[target] == t, m].iloc[0]
                 print(f"Metrics:{t}:{m} {value}")
 
         # Compute and output the confusion matrix
-        conf_mtx = confusion_matrix(y_validate, preds, labels=label_names)
+        conf_mtx = confusion_matrix(y_validate, preds_decoded, labels=label_names)
         for i, row_name in enumerate(label_names):
             for j, col_name in enumerate(label_names):
                 value = conf_mtx[i, j]
@@ -545,9 +587,9 @@ if __name__ == "__main__":
 
     else:
         # Calculate various model performance metrics (regression)
-        rmse = root_mean_squared_error(y_validate, preds)
-        mae = mean_absolute_error(y_validate, preds)
-        r2 = r2_score(y_validate, preds)
+        rmse = root_mean_squared_error(y_validate, preds_decoded)
+        mae = mean_absolute_error(y_validate, preds_decoded)
+        r2 = r2_score(y_validate, preds_decoded)
         print(f"RMSE: {rmse:.3f}")
         print(f"MAE: {mae:.3f}")
         print(f"R2: {r2:.3f}")
@@ -560,7 +602,7 @@ if __name__ == "__main__":
 
     # Save the features (this will validate input during predictions)
     with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
-        json.dump(orig_features, fp)  # We save the original features, not the decompressed ones
+        json.dump(orig_features, fp)
 
     # Save the category mappings
     with open(os.path.join(args.model_dir, "category_mappings.json"), "w") as fp: