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

workbench/model_scripts/pytorch_model/generated_model_script.py

@@ -29,23 +29,20 @@ 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 = {
     "model_type": "regressor",
-    "target": "udm_asy_res_efflux_ratio",
-    "features": ['chi2v', 'fr_sulfone', 'chi1v', 'bcut2d_logplow', 'fr_piperzine', 'kappa3', 'smr_vsa1', 'slogp_vsa5', 'fr_ketone_topliss', 'fr_sulfonamd', 'fr_imine', 'fr_benzene', 'fr_ester', 'chi2n', 'labuteasa', 'peoe_vsa2', 'smr_vsa6', 'bcut2d_chglo', 'fr_sh', 'peoe_vsa1', 'fr_allylic_oxid', 'chi4n', 'fr_ar_oh', 'fr_nh0', 'fr_term_acetylene', 'slogp_vsa7', 'slogp_vsa4', 'estate_vsa1', 'vsa_estate4', 'numbridgeheadatoms', 'numheterocycles', 'fr_ketone', 'fr_morpholine', 'fr_guanido', 'estate_vsa2', 'numheteroatoms', 'fr_nitro_arom_nonortho', 'fr_piperdine', 'nocount', 'numspiroatoms', 'fr_aniline', 'fr_thiophene', 'slogp_vsa10', 'fr_amide', 'slogp_vsa2', 'fr_epoxide', 'vsa_estate7', 'fr_ar_coo', 'fr_imidazole', 'fr_nitrile', 'fr_oxazole', 'numsaturatedrings', 'fr_pyridine', 'fr_hoccn', 'fr_ndealkylation1', 'numaliphaticheterocycles', 'fr_phenol', 'maxpartialcharge', 'vsa_estate5', 'peoe_vsa13', 'minpartialcharge', 'qed', 'fr_al_oh', 'slogp_vsa11', 'chi0n', 'fr_bicyclic', 'peoe_vsa12', 'fpdensitymorgan1', 'fr_oxime', 'molwt', 'fr_dihydropyridine', 'smr_vsa5', 'peoe_vsa5', 'fr_nitro', 'hallkieralpha', 'heavyatommolwt', 'fr_alkyl_halide', 'peoe_vsa8', 'fr_nhpyrrole', 'fr_isocyan', 'bcut2d_chghi', 'fr_lactam', 'peoe_vsa11', 'smr_vsa9', 'tpsa', 'chi4v', 'slogp_vsa1', 'phi', 'bcut2d_logphi', 'avgipc', 'estate_vsa11', 'fr_coo', 'bcut2d_mwhi', 'numunspecifiedatomstereocenters', 'vsa_estate10', 'estate_vsa8', 'numvalenceelectrons', 'fr_nh2', 'fr_lactone', 'vsa_estate1', 'estate_vsa4', 'numatomstereocenters', 'vsa_estate8', 'fr_para_hydroxylation', 'peoe_vsa3', 'fr_thiazole', 'peoe_vsa10', 'fr_ndealkylation2', 'slogp_vsa12', 'peoe_vsa9', 'maxestateindex', 'fr_quatn', 'smr_vsa7', 'minestateindex', 'numaromaticheterocycles', 'numrotatablebonds', 'fr_ar_nh', 'fr_ether', 'exactmolwt', 'fr_phenol_noorthohbond', 'slogp_vsa3', 'fr_ar_n', 'sps', 'fr_c_o_nocoo', 'bertzct', 'peoe_vsa7', 'slogp_vsa8', 'numradicalelectrons', 'molmr', 'fr_tetrazole', 'numsaturatedcarbocycles', 'bcut2d_mrhi', 'kappa1', 'numamidebonds', 'fpdensitymorgan2', 'smr_vsa8', 'chi1n', 'estate_vsa6', 'fr_barbitur', 'fr_diazo', 'kappa2', 'chi0', 'bcut2d_mrlow', 'balabanj', 'peoe_vsa4', 'numhacceptors', 'fr_sulfide', 'chi3n', 'smr_vsa2', 'fr_al_oh_notert', 'fr_benzodiazepine', 'fr_phos_ester', 'fr_aldehyde', 'fr_coo2', 'estate_vsa5', 'fr_prisulfonamd', 'numaromaticcarbocycles', 'fr_unbrch_alkane', 'fr_urea', 'fr_nitroso', 'smr_vsa10', 'fr_c_s', 'smr_vsa3', 'fr_methoxy', 'maxabspartialcharge', 'slogp_vsa9', 'heavyatomcount', 'fr_azide', 'chi3v', 'smr_vsa4', 'mollogp', 'chi0v', 'fr_aryl_methyl', 'fr_nh1', 'fpdensitymorgan3', 'fr_furan', 'fr_hdrzine', 'fr_arn', 'numaromaticrings', 'vsa_estate3', 'fr_azo', 'fr_halogen', 'estate_vsa9', 'fr_hdrzone', 'numhdonors', 'fr_alkyl_carbamate', 'fr_isothiocyan', 'minabspartialcharge', 'fr_al_coo', 'ringcount', 'chi1', 'estate_vsa7', 'fr_nitro_arom', 'vsa_estate9', 'minabsestateindex', 'maxabsestateindex', 'vsa_estate6', 'estate_vsa10', 'estate_vsa3', 'fr_n_o', 'fr_amidine', 'fr_thiocyan', 'fr_phos_acid', 'fr_c_o', 'fr_imide', 'numaliphaticrings', 'peoe_vsa6', 'vsa_estate2', 'nhohcount', 'numsaturatedheterocycles', 'slogp_vsa6', 'peoe_vsa14', 'fractioncsp3', 'bcut2d_mwlow', 'numaliphaticcarbocycles', 'fr_priamide', 'nacid', 'nbase', 'naromatom', 'narombond', 'sz', 'sm', 'sv', 'sse', 'spe', 'sare', 'sp', 'si', 'mz', 'mm', 'mv', 'mse', 'mpe', 'mare', 'mp', 'mi', 'xch_3d', 'xch_4d', 'xch_5d', 'xch_6d', 'xch_7d', 'xch_3dv', 'xch_4dv', 'xch_5dv', 'xch_6dv', 'xch_7dv', 'xc_3d', 'xc_4d', 'xc_5d', 'xc_6d', 'xc_3dv', 'xc_4dv', 'xc_5dv', 'xc_6dv', 'xpc_4d', 'xpc_5d', 'xpc_6d', 'xpc_4dv', 'xpc_5dv', 'xpc_6dv', 'xp_0d', 'xp_1d', 'xp_2d', 'xp_3d', 'xp_4d', 'xp_5d', 'xp_6d', 'xp_7d', 'axp_0d', 'axp_1d', 'axp_2d', 'axp_3d', 'axp_4d', 'axp_5d', 'axp_6d', 'axp_7d', 'xp_0dv', 'xp_1dv', 'xp_2dv', 'xp_3dv', 'xp_4dv', 'xp_5dv', 'xp_6dv', 'xp_7dv', 'axp_0dv', 'axp_1dv', 'axp_2dv', 'axp_3dv', 'axp_4dv', 'axp_5dv', 'axp_6dv', 'axp_7dv', 'c1sp1', 'c2sp1', 'c1sp2', 'c2sp2', 'c3sp2', 'c1sp3', 'c2sp3', 'c3sp3', 'c4sp3', 'hybratio', 'fcsp3', 'num_stereocenters', 'num_unspecified_stereocenters', 'num_defined_stereocenters', 'num_r_centers', 'num_s_centers', 'num_stereobonds', 'num_e_bonds', 'num_z_bonds', 'stereo_complexity', 'frac_defined_stereo'],
+    "target": "caco_2_efflux",
+    "features": ['smr_vsa3', 'xch_6dv', 'tpsa', 'numhdonors', 'vsa_estate2', 'fr_imidazole', 'mollogp', 'molmr', 'vsa_estate3', 'bcut2d_mwhi', 'slogp_vsa1', 'peoe_vsa6', 'peoe_vsa9', 'fr_halogen', 'vsa_estate10', 'fr_pyridine', 'minabsestateindex', 'bcut2d_mwlow', 'peoe_vsa2', 'bcut2d_logplow', 'vsa_estate7', 'peoe_vsa1', 'bcut2d_mrlow', 'slogp_vsa4', 'peoe_vsa7', 'peoe_vsa8', 'bcut2d_logphi', 'estate_vsa8', 'sv', 'fpdensitymorgan1', 'peoe_vsa3', 'slogp_vsa2', 'vsa_estate8', 'hallkieralpha', 'chi2v', 'axp_3dv', 'estate_vsa7', 'peoe_vsa11', 'xc_3d', 'vsa_estate4', 'vsa_estate9', 'bcut2d_chghi', 'vsa_estate5', 'c1sp2', 'slogp_vsa6', 'axp_3d', 'axp_2dv', 'bcut2d_mrhi', 'xc_3dv', 'xpc_4d', 'chi3n', 'xp_3dv', 'kappa3', 'vsa_estate6', 'minestateindex', 'mp', 'avgipc', 'axp_5dv', 'heavyatommolwt', 'maxpartialcharge', 'peoe_vsa10', 'kappa1', 'nocount', 'xp_6d', 'xpc_4dv', 'axp_4dv', 'sz', 'axp_4d', 'xc_5d', 'qed', 'xch_7dv', 'axp_1dv', 'chi0v', 'smr_vsa5', 'maxabspartialcharge', 'minpartialcharge', 'estate_vsa5', 'fpdensitymorgan3', 'fpdensitymorgan2', 'labuteasa', 'fr_methoxy', 'kappa2', 'sps', 'mi', 'bcut2d_chglo', 'vsa_estate1', 'fr_nh2', 'axp_1d', 'maxestateindex', 'estate_vsa3', 'bertzct', 'phi', 'fractioncsp3', 'xc_5dv', 'smr_vsa6', 'estate_vsa6', 'fr_ar_coo', 'axp_7dv', 'slogp_vsa8', 'estate_vsa2', 'axp_2d', 'balabanj', 'mz', 'estate_vsa4', 'smr_vsa7', 'slogp_vsa5', 'chi0n', 'xp_5d', 'xch_7d', 'chi4n', 'smr_vsa10', 'numheterocycles', 'chi1v', 'axp_5d', 'axp_0dv', 'fr_al_oh', 'chi2n', 'fr_nh0', 'chi1n', 'xp_7dv', 'xpc_5d', 'ringcount', 'xpc_5dv', 'peoe_vsa12', 'xpc_6dv', 'xp_6dv', 'mse', 'sp', 'xpc_6d', 'chi4v', 'fr_aniline', 'c2sp3', 'peoe_vsa4', 'numheteroatoms', 'smr_vsa9', 'xp_7d', 'axp_6d', 'molwt', 'numrotatablebonds', 'smr_vsa4', 'stereo_complexity', 'axp_7d', 'slogp_vsa3', 'spe', 'num_r_centers', 'si', 'axp_0d', 'xp_2d', 'xp_4d', 'c2sp2', 'fr_aryl_methyl', 'mpe', 'xch_6d', 'axp_6dv', 'numsaturatedcarbocycles', 'fr_para_hydroxylation', 'estate_vsa10', 'estate_vsa1', 'hybratio', 'numhacceptors', 'naromatom', 'chi1', 'fr_urea', 'xp_3d', 'smr_vsa1', 'num_s_centers', 'xch_5dv', 'c3sp2', 'mare', 'xp_5dv', 'fr_al_oh_notert', 'estate_vsa9', 'fr_piperdine', 'numunspecifiedatomstereocenters', 'chi3v', 'c3sp3', 'chi0', 'numsaturatedheterocycles', 'xp_4dv', 'fr_amide', 'fr_nhpyrrole', 'mv', 'fr_ar_n', 'xc_4dv', 'fr_morpholine', 'fr_ndealkylation2', 'xch_3d', 'xch_4d', 'slogp_vsa10', 'fr_ar_oh', 'fr_benzene', 'fr_nh1', 'c1sp1', 'sse', 'num_defined_stereocenters', 'xch_4dv', 'peoe_vsa14', 'xch_5d', 'fr_hoccn', 'fr_nitrile', 'mm', 'fr_priamide', 'xc_6dv', 'num_unspecified_stereocenters', 'fr_ether', 'fr_piperzine', 'fr_bicyclic', 'fr_term_acetylene'],
     "compressed_features": [],
-    "model_metrics_s3_path": "s3://ideaya-sageworks-bucket/models/caco2-er-reg-pytorch-test/training",
-    "train_all_data": False,
+    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/caco2-efflux-ref-pytorch/training",
+    "train_all_data": True,
     "hyperparameters": {},
 }
 
 
-# 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(
@@ -529,7 +571,7 @@ if __name__ == "__main__":
             }
         )
 
-        # We need to get creative with the Classification Metrics
+        # Output metrics per class
         metrics = ["precision", "recall", "f1", "support"]
         for t in label_names:
             for m in metrics:
@@ -537,7 +579,7 @@ if __name__ == "__main__":
                 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: