workbench 0.8.156__py3-none-any.whl → 0.8.158__py3-none-any.whl

{workbench-0.8.156.dist-info → workbench-0.8.158.dist-info}/RECORD

@@ -51,16 +51,16 @@ workbench/core/artifacts/__init__.py,sha256=ps7rA_rbWnDbvWbg4kvu--IKMY8WmbPRyv4S
 workbench/core/artifacts/artifact.py,sha256=mF1lqQ8EP43zMRSxqPYuWoHRwQTjyWdpW9LewU7utEE,17517
 workbench/core/artifacts/athena_source.py,sha256=RNmCe7s6uH4gVHpcdJcL84aSbF5Q1ahJBLLGwHYRXEU,26081
 workbench/core/artifacts/cached_artifact_mixin.py,sha256=TNRq-4frDS0GRodfnz6qN2SGP0U_QglZw5X5gw-K7-Y,3755
-workbench/core/artifacts/data_source_abstract.py,sha256=27WVCQkIDO-cj0URdDu-9LLnuWxU_3_0wyLt4iFIwsg,10604
+workbench/core/artifacts/data_source_abstract.py,sha256=5IRCzFVK-17cd4NXPMRfx99vQAmQ0WHE5jcm5RfsVTg,10619
 workbench/core/artifacts/data_source_factory.py,sha256=YL_tA5fsgubbB3dPF6T4tO0rGgz-6oo3ge4i_YXVC-M,2380
 workbench/core/artifacts/endpoint_core.py,sha256=L6uWOxHKItjbpRS2rFrAbxAqDyZIv2CO9dnZpohKrUI,48768
-workbench/core/artifacts/feature_set_core.py,sha256=7yyy_0RFmrhcXqZHVb0P6suo3ye1tSO_kJqqiGgjrFw,29377
+workbench/core/artifacts/feature_set_core.py,sha256=055VdSYR09HP4ygAuYvIYtHQ7Ec4XxsZygpgEl5H5jQ,29136
 workbench/core/artifacts/model_core.py,sha256=WjiZyqwUVeVZF22FO8zXFoHmgh-Vd6c_KQLx_da_FTQ,50194
 workbench/core/artifacts/monitor_core.py,sha256=BP6UuCyBI4zB2wwcIXvUw6RC0EktTcQd5Rv0x73qzio,37670
 workbench/core/cloud_platform/cloud_meta.py,sha256=J3qqMLdzikRRI9TC7UEk9Ib2-64ovU8X-lzLxhze1uY,8669
 workbench/core/cloud_platform/aws/README.md,sha256=QT5IQXoUHbIA0qQ2wO6_2P2lYjYQFVYuezc22mWY4i8,97
-workbench/core/cloud_platform/aws/aws_account_clamp.py,sha256=YRZ_DufSNfuKgW7q9oanSwqlhZecbnapNW-twBajUX8,5683
-workbench/core/cloud_platform/aws/aws_df_store.py,sha256=Tdgg8QThmjWQs5gXgjK7D34qaQKe7MsaO5s8fD133F0,15043
+workbench/core/cloud_platform/aws/aws_account_clamp.py,sha256=OzFknZXKW7VTvnDGGX4BXKoh0i1gQ7yaEBhkLCyHFSs,6310
+workbench/core/cloud_platform/aws/aws_df_store.py,sha256=utRIlTCPwFneHHZ8_Z3Hw3rOJSeryiFA4wBtucxULRQ,15055
 workbench/core/cloud_platform/aws/aws_graph_store.py,sha256=ytYxQTplUmeWbsPmxyZbf6mO9qyTl60ewlJG8MyfyEY,9414
 workbench/core/cloud_platform/aws/aws_meta.py,sha256=ijbzucUF948RyD-GgYYyxl-Ejvkv5Z6gqB_g497ss7M,33932
 workbench/core/cloud_platform/aws/aws_parameter_store.py,sha256=9ekuMOQFHFMIEV68UbHhS_fLB9iqG5Hvu4EV6iamEpk,10400
@@ -101,14 +101,14 @@ workbench/core/transforms/features_to_features/__init__.py,sha256=47DEQpj8HBSa-_
 workbench/core/transforms/features_to_features/heavy/emr/Readme.md,sha256=YtQgCEQeKe0CQXQkhzMTYq9xOtCsCYb5P5LW2BmRKWQ,68
 workbench/core/transforms/features_to_features/heavy/glue/Readme.md,sha256=TuyCatWfoDr99zUwvOcxf-TqMkQzaMqXlj5nmFcRzfo,48
 workbench/core/transforms/features_to_model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-workbench/core/transforms/features_to_model/features_to_model.py,sha256=qtqfiN4a35TqxP_Sxfu90zrbNMA3r-q0F13Yr0VaIbM,19162
+workbench/core/transforms/features_to_model/features_to_model.py,sha256=xAOcbcm2f9iumCCM_T8OTdzRkl3ZLTzyGRdv9ImmnqM,19333
 workbench/core/transforms/model_to_endpoint/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workbench/core/transforms/model_to_endpoint/model_to_endpoint.py,sha256=5IBhA56qSDSnfh4Xwvb2PP5UDM7md2R7Ur38hP4Mgyw,4624
 workbench/core/transforms/pandas_transforms/__init__.py,sha256=xL4MT8-fZ1SFqDbTLc8XyxjupHtB1YR6Ej0AC2nwd7I,894
 workbench/core/transforms/pandas_transforms/data_to_pandas.py,sha256=sJHPeuNF8Q8aQqgRnkdWkyvur5cbggdUVIwR-xF3Dlo,3621
 workbench/core/transforms/pandas_transforms/features_to_pandas.py,sha256=af6xdPt2V4zhh-SzQa_UYxdmNMzMLXbrbsznV5QoIJg,3441
 workbench/core/transforms/pandas_transforms/pandas_to_data.py,sha256=cqo6hQmzUGUFACvNuVLZQdgrlXrQIu4NjqK-ujPmoIc,9181
-workbench/core/transforms/pandas_transforms/pandas_to_features.py,sha256=wdPFnnVUBbMK2_xTSbGJ9x2TESHf6znHl3LwXtPxlX0,20183
+workbench/core/transforms/pandas_transforms/pandas_to_features.py,sha256=XiUz2BqOX4N34g6hvFvcLswhkEouyU0AjyIy9EGv2zg,20440
 workbench/core/transforms/pandas_transforms/pandas_to_features_chunked.py,sha256=0R8mQlWfbIlTVmYUmrtu2gsw0AE815k6kqPgpd0bmyQ,4422
 workbench/core/views/__init__.py,sha256=UZJMAJBCMVM3uSYmnFg8c2LWtdu9-479WNAdVMIohAc,962
 workbench/core/views/column_subset_view.py,sha256=vGDKTTGrPIY-IFOeWvudJrhKiq0OjWDp5rTuuj-X40U,4261
@@ -130,6 +130,10 @@ workbench/model_scripts/custom_models/chem_info/requirements.txt,sha256=7HBUzvNi
 workbench/model_scripts/custom_models/chem_info/tautomerize.py,sha256=KAxTAqtTql4_FvnrAyYRgaJEmtAx399HXA_iw_awa08,3125
 workbench/model_scripts/custom_models/meta_endpoints/example.py,sha256=hzOAuLhIGB8vei-555ruNxpsE1GhuByHGjGB0zw8GSs,1726
 workbench/model_scripts/custom_models/network_security/Readme.md,sha256=Z2gtiu0hLHvEJ1x-_oFq3qJZcsK81sceBAGAGltpqQ8,222
+workbench/model_scripts/custom_models/nn_models/Readme.md,sha256=x6U_gox2yV-kLspgmzE77t2xk5GFNgcpcuQq-Q78G7w,146
+workbench/model_scripts/custom_models/nn_models/generated_model_script.py,sha256=PsowW-kosyN2wWhuDgP_41mBYscUDYKOhzzQFVUKBzc,20695
+workbench/model_scripts/custom_models/nn_models/pytorch.template,sha256=-gd0FbDW1ilTo1WcwiOLVZPAyTQ09naVlnSXxrVxKYU,20422
+workbench/model_scripts/custom_models/nn_models/requirements.txt,sha256=sC6v1LSBkwJFbvObn6DUD1HwPM86_rbmXEMH5Tcn2kM,184
 workbench/model_scripts/custom_models/proximity/Readme.md,sha256=RlMFAJZgAT2mCgDk-UwR_R0Y_NbCqeI5-8DUsxsbpWQ,289
 workbench/model_scripts/custom_models/proximity/feature_space_proximity.template,sha256=2c3eDu4sLP_bCTBAf_aIR1QdC7CpYDpXo8UU_2ZoLuE,4833
 workbench/model_scripts/custom_models/proximity/generated_model_script.py,sha256=RdbKbXtrSNYQJvB-oLcRHpJ6w0TM7zbmMfuocHb7GM0,7967
@@ -142,7 +146,6 @@ workbench/model_scripts/custom_models/uq_models/gaussian_process.template,sha256
 workbench/model_scripts/custom_models/uq_models/generated_model_script.py,sha256=l74VibzFnhmPeNUEiFwIIg5aNujcCs9LtRywUvc5Avo,14528
 workbench/model_scripts/custom_models/uq_models/mapie_xgb.template,sha256=ZTmerwkmXtewJwx3GGJSdLRyzJV5SJ86PvCu3dV_GHw,7330
 workbench/model_scripts/custom_models/uq_models/meta_uq.template,sha256=26FNangcpyV9nFOIufRuVZ45BQv6oPf9xlJZkVIULG4,9287
-workbench/model_scripts/custom_models/uq_models/meta_uq_experiment.template,sha256=2MT6-jzhmz69DuQCXAmMXoKRkBruto1m92LDsG8vdVI,14326
 workbench/model_scripts/custom_models/uq_models/ngboost.template,sha256=N-eWP967-X2Qbvk18VL7LPXRJMKne9SS2fb_jntwTec,7738
 workbench/model_scripts/custom_models/uq_models/proximity.py,sha256=zqmNlX70LnWXr5fdtFFQppSNTLjlOciQVrjGr-g9jRE,13716
 workbench/model_scripts/custom_models/uq_models/requirements.txt,sha256=jfwV5b1t6BFtdaRGrSz8LnuQzJm-4V5OlhhP-4CGxhs,107
@@ -155,7 +158,7 @@ workbench/model_scripts/quant_regression/quant_regression.template,sha256=AQihff
 workbench/model_scripts/quant_regression/requirements.txt,sha256=jWlGc7HH7vqyukTm38LN4EyDi8jDUPEay4n45z-30uc,104
 workbench/model_scripts/scikit_learn/requirements.txt,sha256=aVvwiJ3LgBUhM_PyFlb2gHXu_kpGPho3ANBzlOkfcvs,107
 workbench/model_scripts/scikit_learn/scikit_learn.template,sha256=d4pgeZYFezUQsB-7iIsjsUgB1FM6d27651wpfDdXmI0,12640
-workbench/model_scripts/xgb_model/generated_model_script.py,sha256=aqrDobLcMysjc-q8p5aCyPXPrVlia14CbARRqaG8--Q,21083
+workbench/model_scripts/xgb_model/generated_model_script.py,sha256=BDExjY02zSbcCloW4VNzHvej_kS3H0uJxa92egoASNY,17954
 workbench/model_scripts/xgb_model/requirements.txt,sha256=jWlGc7HH7vqyukTm38LN4EyDi8jDUPEay4n45z-30uc,104
 workbench/model_scripts/xgb_model/xgb_model.template,sha256=RaUr8X6al5R2IILNKgGUH05Gb4H7AFFG9RE524_VH7Q,17935
 workbench/repl/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -194,7 +197,7 @@ workbench/utils/aws_utils.py,sha256=XckM0vzud7Nx1OxD1GoYGLQxdj1PqeQ43cN66tnrRYI,
 workbench/utils/bulk_utils.py,sha256=s1lYN2Uk536MNGetekLYL_VL0N34hUjk1FX9BAz3Qu0,1182
 workbench/utils/cache.py,sha256=0R5RXYEz_XHARK3anmQC4VRMawMks_cJ8S4vwC2roAE,5524
 workbench/utils/chem_utils.py,sha256=tLTAvLKTOiYSzbVQF0M8V5-ej36IVgr21CNB2vVJjYQ,56780
-workbench/utils/cloudwatch_handler.py,sha256=ocrj4oXIBIBZHax_ekk4E4r-tD-yTrlDmZeKRFBauOo,5491
+workbench/utils/cloudwatch_handler.py,sha256=dtnkr8tXtTRAASQ60QO0lz3SRA5LEbzsK1VCIqblfKs,5157
 workbench/utils/color_utils.py,sha256=TmDGLK44t975lkfjt_1O-ee02QxrKfke7vPuXb-V-Uo,11779
 workbench/utils/config_manager.py,sha256=Yj43Ta67dn34XdOcKcOvXw38ln6TRiv4DePXjPt2jg8,17641
 workbench/utils/dashboard_metrics.py,sha256=cNFI0GIAjd_IiDzM1oebsJ2QkRZuW068W_66ZC3J100,7398
@@ -252,7 +255,7 @@ workbench/web_interface/components/experiments/dashboard_metric_plots.py,sha256=
 workbench/web_interface/components/experiments/outlier_plot.py,sha256=5bWsmJEXyt50npeQxLHXCPtiq4WRVgg938Sl0DVjNWg,3647
 workbench/web_interface/components/plugins/ag_table.py,sha256=HrPOMotlOGigk0v8Cxx_doSHXdOKTT1-bzlsqDwwzng,3979
 workbench/web_interface/components/plugins/confusion_matrix.py,sha256=1K94JSlDwQwdf5uDYVydQzY-EQm89hYXchxbXoNvons,7176
-workbench/web_interface/components/plugins/dashboard_status.py,sha256=KWQW6EArOXhrUwZ0mMZydQk6nOpIoU6ftP4a0goN-Ew,5690
+workbench/web_interface/components/plugins/dashboard_status.py,sha256=8Tu38lR5YgntxDjz_x2XfLiW7SOdreNLOFT5VkbYzKo,5748
 workbench/web_interface/components/plugins/data_details.py,sha256=pZm1AbM_0EXQwx77qUkfyrU9MedAs4Wlkp6iOtSrUtI,11104
 workbench/web_interface/components/plugins/endpoint_details.py,sha256=0A7g_Lx5-3XnDWOGT3YEDPNpmME_-WfYc65f-rRVjJE,3769
 workbench/web_interface/components/plugins/generated_compounds.py,sha256=hC0sh-1_rbN55Huno-E_2wF37kgIHi5Mtaer6Xk5fRM,8052
@@ -272,9 +275,9 @@ workbench/web_interface/page_views/main_page.py,sha256=QIGiQPXu5Q9TU15nPEAYa6sjQ
 workbench/web_interface/page_views/models_page_view.py,sha256=M0bdC7bAzLyIaE2jviY12FF4abdMFZmg6sFuOY_LaGI,2650
 workbench/web_interface/page_views/page_view.py,sha256=Gh6YnpOGlUejx-bHZAf5pzqoQ1H1R0OSwOpGhOBO06w,455
 workbench/web_interface/page_views/pipelines_page_view.py,sha256=v2pxrIbsHBcYiblfius3JK766NZ7ciD2yPx0t3E5IJo,2656
-workbench-0.8.156.dist-info/licenses/LICENSE,sha256=z4QMMPlLJkZjU8VOKqJkZiQZCEZ--saIU2Z8-p3aVc0,1080
-workbench-0.8.156.dist-info/METADATA,sha256=29854fjJPyoP6gHfmAt9DQFAI_KuS72Cjh2q133azrk,9497
-workbench-0.8.156.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-workbench-0.8.156.dist-info/entry_points.txt,sha256=oZykkheWiiIBjRE8cS5SdcxwmZKSFaQEGwMBjNh-eNM,238
-workbench-0.8.156.dist-info/top_level.txt,sha256=Dhy72zTxaA_o_yRkPZx5zw-fwumnjGaeGf0hBN3jc_w,10
-workbench-0.8.156.dist-info/RECORD,,
+workbench-0.8.158.dist-info/licenses/LICENSE,sha256=z4QMMPlLJkZjU8VOKqJkZiQZCEZ--saIU2Z8-p3aVc0,1080
+workbench-0.8.158.dist-info/METADATA,sha256=gzWmV7xJLnwpDa2r_SM6XIOPXR2zNb6shhtoEwHD1d4,9393
+workbench-0.8.158.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+workbench-0.8.158.dist-info/entry_points.txt,sha256=oZykkheWiiIBjRE8cS5SdcxwmZKSFaQEGwMBjNh-eNM,238
+workbench-0.8.158.dist-info/top_level.txt,sha256=Dhy72zTxaA_o_yRkPZx5zw-fwumnjGaeGf0hBN3jc_w,10
+workbench-0.8.158.dist-info/RECORD,,

workbench/model_scripts/custom_models/uq_models/meta_uq_experiment.template

@@ -1,393 +0,0 @@
-# Model: NGBoost Regressor with Distribution output
-from ngboost import NGBRegressor
-from xgboost import XGBRegressor  # Base Estimator
-from sklearn.model_selection import train_test_split
-import numpy as np
-
-# Model Performance Scores
-from sklearn.metrics import (
-    mean_absolute_error,
-    r2_score,
-    root_mean_squared_error
-)
-
-from io import StringIO
-import json
-import argparse
-import joblib
-import os
-import pandas as pd
-
-# Local Imports
-from proximity import Proximity
-
-
-
-# Template Placeholders
-TEMPLATE_PARAMS = {
-    "id_column": "{{id_column}}",
-    "features": "{{feature_list}}",
-    "target": "{{target_column}}",
-    "train_all_data": "{{train_all_data}}",
-    "track_columns": "{{track_columns}}"
-}
-
-
-# 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.
-
-    Args:
-        df (pd.DataFrame): DataFrame to check
-        df_name (str): Name of the DataFrame
-    """
-    if df.empty:
-        msg = f"*** The training data {df_name} has 0 rows! ***STOPPING***"
-        print(msg)
-        raise ValueError(msg)
-
-
-def match_features_case_insensitive(df: pd.DataFrame, model_features: list) -> pd.DataFrame:
-    """
-    Matches and renames DataFrame columns to match model feature names (case-insensitive).
-    Prioritizes exact matches, then case-insensitive matches.
-
-    Raises ValueError if any model features cannot be matched.
-    """
-    df_columns_lower = {col.lower(): col for col in df.columns}
-    rename_dict = {}
-    missing = []
-    for feature in model_features:
-        if feature in df.columns:
-            continue  # Exact match
-        elif feature.lower() in df_columns_lower:
-            rename_dict[df_columns_lower[feature.lower()]] = feature
-        else:
-            missing.append(feature)
-
-    if missing:
-        raise ValueError(f"Features not found: {missing}")
-
-    # Rename the DataFrame columns to match the model features
-    return df.rename(columns=rename_dict)
-
-
-def distance_weighted_calibrated_intervals(
-        df_pred: pd.DataFrame,
-        prox_df: pd.DataFrame,
-        calibration_strength: float = 0.7,
-        distance_decay: float = 3.0,
-) -> pd.DataFrame:
-    """
-    Calibrate intervals using distance-weighted neighbor quantiles.
-    Uses all 10 neighbors with distance-based weighting.
-    """
-    id_column = TEMPLATE_PARAMS["id_column"]
-    target_column = TEMPLATE_PARAMS["target"]
-
-    # Distance-weighted neighbor statistics
-    def weighted_quantile(values, weights, q):
-        """Calculate weighted quantile"""
-        if len(values) == 0:
-            return np.nan
-        sorted_indices = np.argsort(values)
-        sorted_values = values[sorted_indices]
-        sorted_weights = weights[sorted_indices]
-        cumsum = np.cumsum(sorted_weights)
-        cutoff = q * cumsum[-1]
-        return np.interp(cutoff, cumsum, sorted_values)
-
-    # Calculate distance weights (closer neighbors get more weight)
-    prox_df = prox_df.copy()
-    prox_df['weight'] = 1 / (1 + prox_df['distance'] ** distance_decay)
-
-    # Get weighted quantiles and statistics for each ID
-    neighbor_stats = []
-    for id_val, group in prox_df.groupby(id_column):
-        values = group[target_column].values
-        weights = group['weight'].values
-
-        # Normalize weights
-        weights = weights / weights.sum()
-
-        stats = {
-            id_column: id_val,
-            'local_q025': weighted_quantile(values, weights, 0.025),
-            'local_q25': weighted_quantile(values, weights, 0.25),
-            'local_q75': weighted_quantile(values, weights, 0.75),
-            'local_q975': weighted_quantile(values, weights, 0.975),
-            'local_median': weighted_quantile(values, weights, 0.5),
-            'local_std': np.sqrt(np.average((values - np.average(values, weights=weights)) ** 2, weights=weights)),
-            'avg_distance': group['distance'].mean(),
-            'min_distance': group['distance'].min(),
-            'max_distance': group['distance'].max(),
-        }
-        neighbor_stats.append(stats)
-
-    neighbor_df = pd.DataFrame(neighbor_stats)
-    out = df_pred.merge(neighbor_df, on=id_column, how='left')
-
-    # Model disagreement score (normalized by prediction std)
-    model_disagreement = (out["prediction"] - out["prediction_uq"]).abs()
-    disagreement_score = (model_disagreement / out["prediction_std"]).clip(0, 2)
-
-    # Local confidence based on:
-    # 1. How close the neighbors are (closer = more confident)
-    # 2. How much local variance there is (less variance = more confident)
-    max_reasonable_distance = out['max_distance'].quantile(0.8)  # 80th percentile as reference
-    distance_confidence = (1 - (out['avg_distance'] / max_reasonable_distance)).clip(0.1, 1.0)
-
-    variance_confidence = (out["prediction_std"] / out["local_std"]).clip(0.5, 2.0)
-    local_confidence = distance_confidence * variance_confidence.clip(0.5, 1.5)
-
-    # Calibration weight: higher when models disagree and we have good local data
-    calibration_weight = (
-            calibration_strength *
-            local_confidence *  # Weight by local data quality
-            disagreement_score.clip(0.3, 1.0)  # More calibration when models disagree
-    )
-
-    # Consensus prediction (slight preference for NGBoost since it provides intervals)
-    consensus_pred = 0.65 * out["prediction_uq"] + 0.35 * out["prediction"]
-
-    # Re-center local intervals around consensus prediction
-    local_center_offset = consensus_pred - out["local_median"]
-
-    # Apply calibration to each quantile
-    quantile_pairs = [
-        ("q_025", "local_q025"),
-        ("q_25", "local_q25"),
-        ("q_75", "local_q75"),
-        ("q_975", "local_q975")
-    ]
-
-    for model_q, local_q in quantile_pairs:
-        # Adjust local quantiles to be centered around consensus
-        adjusted_local_q = out[local_q] + local_center_offset
-
-        # Blend model and local intervals
-        out[model_q] = (
-                (1 - calibration_weight) * out[model_q] +
-                calibration_weight * adjusted_local_q
-        )
-
-    # Ensure proper interval ordering and bounds using pandas
-    out["q_025"] = pd.concat([out["q_025"], consensus_pred], axis=1).min(axis=1)
-    out["q_975"] = pd.concat([out["q_975"], consensus_pred], axis=1).max(axis=1)
-    out["q_25"] = pd.concat([out["q_25"], out["q_75"]], axis=1).min(axis=1)
-
-    # Optional: Add some interval expansion when neighbors are very far
-    # (indicates we're in a sparse region of feature space)
-    sparse_region_mask = out['min_distance'] > out['min_distance'].quantile(0.9)
-    expansion_factor = 1 + 0.2 * sparse_region_mask  # 20% expansion in sparse regions
-
-    for q in ["q_025", "q_25", "q_75", "q_975"]:
-        interval_width = out[q] - consensus_pred
-        out[q] = consensus_pred + interval_width * expansion_factor
-
-    # Clean up temporary columns
-    cleanup_cols = [col for col in out.columns if col.startswith("local_")] + \
-                   ['avg_distance', 'min_distance', 'max_distance']
-
-    return out.drop(columns=cleanup_cols)
-
-
-# TRAINING SECTION
-#
-# This section (__main__) is where SageMaker will execute the training job
-# and save the model artifacts to the model directory.
-#
-if __name__ == "__main__":
-    # Template Parameters
-    id_column = TEMPLATE_PARAMS["id_column"]
-    features = TEMPLATE_PARAMS["features"]
-    target = TEMPLATE_PARAMS["target"]
-    train_all_data = TEMPLATE_PARAMS["train_all_data"]
-    track_columns = TEMPLATE_PARAMS["track_columns"]  # Can be None
-    validation_split = 0.2
-
-    # Script arguments for input/output directories
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
-    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
-    parser.add_argument(
-        "--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data")
-    )
-    args = parser.parse_args()
-
-    # Load training data from the specified directory
-    training_files = [
-        os.path.join(args.train, file)
-        for file in os.listdir(args.train) if file.endswith(".csv")
-    ]
-    print(f"Training Files: {training_files}")
-
-    # Combine files and read them all into a single pandas dataframe
-    df = pd.concat([pd.read_csv(file, engine="python") for file in training_files])
-
-    # Check if the DataFrame is empty
-    check_dataframe(df, "training_df")
-
-    # Training data split logic
-    if train_all_data:
-        # Use all data for both training and validation
-        print("Training on all data...")
-        df_train = df.copy()
-        df_val = df.copy()
-    elif "training" in df.columns:
-        # Split data based on a 'training' column if it exists
-        print("Splitting data based on 'training' column...")
-        df_train = df[df["training"]].copy()
-        df_val = df[~df["training"]].copy()
-    else:
-        # Perform a random split if no 'training' column is found
-        print("Splitting data randomly...")
-        df_train, df_val = train_test_split(df, test_size=validation_split, random_state=42)
-
-    # We're using XGBoost for point predictions and NGBoost for uncertainty quantification
-    xgb_model = XGBRegressor()
-    ngb_model = NGBRegressor()
-
-    # Prepare features and targets for training
-    X_train = df_train[features]
-    X_val = df_val[features]
-    y_train = df_train[target]
-    y_val = df_val[target]
-
-    # Train both models using the training data
-    xgb_model.fit(X_train, y_train)
-    ngb_model.fit(X_train, y_train, X_val=X_val, Y_val=y_val)
-
-    # Make Predictions on the Validation Set
-    print(f"Making Predictions on Validation Set...")
-    y_validate = df_val[target]
-    X_validate = df_val[features]
-    preds = xgb_model.predict(X_validate)
-
-    # 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)
-    print(f"RMSE: {rmse:.3f}")
-    print(f"MAE: {mae:.3f}")
-    print(f"R2: {r2:.3f}")
-    print(f"NumRows: {len(df_val)}")
-
-    # Save the trained XGBoost model
-    xgb_model.save_model(os.path.join(args.model_dir, "xgb_model.json"))
-
-    # Save the trained NGBoost model
-    joblib.dump(ngb_model, os.path.join(args.model_dir, "ngb_model.joblib"))
-
-    # Save the feature list to validate input during predictions
-    with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
-        json.dump(features, fp)
-
-    # Now the Proximity model
-    model = Proximity(df_train, id_column, features, target, track_columns=track_columns)
-
-    # Now serialize the model
-    model.serialize(args.model_dir)
-
-
-#
-# Inference Section
-#
-def model_fn(model_dir) -> dict:
-    """Load and return XGBoost and NGBoost regressors from model directory."""
-
-    # Load XGBoost regressor
-    xgb_path = os.path.join(model_dir, "xgb_model.json")
-    xgb_model = XGBRegressor(enable_categorical=True)
-    xgb_model.load_model(xgb_path)
-
-    # Load NGBoost regressor
-    ngb_model = joblib.load(os.path.join(model_dir, "ngb_model.joblib"))
-
-    # Deserialize the proximity model
-    prox_model = Proximity.deserialize(model_dir)
-
-    return {
-        "xgboost": xgb_model,
-        "ngboost": ngb_model,
-        "proximity": prox_model
-    }
-
-
-def input_fn(input_data, content_type):
-    """Parse input data and return a DataFrame."""
-    if not input_data:
-        raise ValueError("Empty input data is not supported!")
-
-    # Decode bytes to string if necessary
-    if isinstance(input_data, bytes):
-        input_data = input_data.decode("utf-8")
-
-    if "text/csv" in content_type:
-        return pd.read_csv(StringIO(input_data))
-    elif "application/json" in content_type:
-        return pd.DataFrame(json.loads(input_data))  # Assumes JSON array of records
-    else:
-        raise ValueError(f"{content_type} not supported!")
-
-
-def output_fn(output_df, accept_type):
-    """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
-        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)
-    else:
-        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
-
-
-def predict_fn(df, models) -> pd.DataFrame:
-    """Make Predictions with our XGB Quantile Regression Model
-
-    Args:
-        df (pd.DataFrame): The input DataFrame
-        models (dict): The dictionary of models to use for predictions
-
-    Returns:
-        pd.DataFrame: The DataFrame with the predictions added
-    """
-
-    # Grab our feature columns (from training)
-    model_dir = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
-    with open(os.path.join(model_dir, "feature_columns.json")) as fp:
-        model_features = json.load(fp)
-
-    # Match features in a case-insensitive manner
-    matched_df = match_features_case_insensitive(df, model_features)
-
-    # Use XGBoost for point predictions
-    df["prediction"] = models["xgboost"].predict(matched_df[model_features])
-
-    # NGBoost predict returns distribution objects
-    y_dists = models["ngboost"].pred_dist(matched_df[model_features])
-
-    # Extract parameters from distribution
-    dist_params = y_dists.params
-
-    # Extract mean and std from distribution parameters
-    df["prediction_uq"] = dist_params['loc']  # mean
-    df["prediction_std"] = dist_params['scale']  # standard deviation
-
-    # Add 95% prediction intervals using ppf (percent point function)
-    df["q_025"] = y_dists.ppf(0.025)  # 2.5th percentile
-    df["q_975"] = y_dists.ppf(0.975)  # 97.5th percentile
-
-    # Add 50% prediction intervals
-    df["q_25"] = y_dists.ppf(0.25)   # 25th percentile
-    df["q_75"] = y_dists.ppf(0.75)   # 75th percentile
-
-    # Compute Nearest neighbors with Proximity model
-    prox_df = models["proximity"].neighbors(df)
-
-    # Shrink prediction intervals based on KNN variance
-    df = distance_weighted_calibrated_intervals(df, prox_df)
-
-    # Return the modified DataFrame
-    return df