File size: 5,342 Bytes
2532605 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """
KRONECTOR - LightGBM training entry point.
Training flow:
- load race parquet data
- prepare features and fit categorical LabelEncoders
- evaluate all TimeSeriesSplit folds and log averaged CV metrics
- retrain a final LightGBM model on the full dataset
- log the model, fitted encoders, and SHAP summary to MLflow
"""
from __future__ import annotations
import argparse
import os
import tempfile
from pathlib import Path
import numpy as np
import pandas as pd
from sklearn.metrics import log_loss, roc_auc_score
from ml.feature_engineering import (
create_time_series_splits,
prepare_model_data,
save_encoders,
)
REGISTERED_MODEL_NAME = "kronector-f1-lgbm"
LIGHTGBM_PARAMS = {
"objective": "binary",
"n_estimators": 200,
"learning_rate": 0.05,
"num_leaves": 31,
"random_state": 42,
"class_weight": "balanced",
"verbosity": -1,
}
def _positive_class_shap_values(shap_values):
"""Normalize SHAP binary-class outputs to one 2D array."""
if isinstance(shap_values, list):
return shap_values[1] if len(shap_values) > 1 else shap_values[0]
if isinstance(shap_values, np.ndarray) and shap_values.ndim == 3:
return shap_values[:, :, 1]
return shap_values
def _build_model():
import lightgbm as lgb
return lgb.LGBMClassifier(**LIGHTGBM_PARAMS)
def _cross_validate(bundle, n_splits: int) -> dict[str, float]:
"""Evaluate all time-series folds and return averaged metrics."""
fold_metrics = []
for fold, (train_idx, valid_idx) in enumerate(
create_time_series_splits(bundle.X, n_splits=n_splits), start=1
):
model = _build_model()
X_train = bundle.X.iloc[train_idx]
y_train = bundle.y.iloc[train_idx]
X_valid = bundle.X.iloc[valid_idx]
y_valid = bundle.y.iloc[valid_idx]
model.fit(X_train, y_train)
valid_prob = model.predict_proba(X_valid)[:, 1]
metrics = {
"fold": fold,
"log_loss": log_loss(y_valid, valid_prob, labels=[0, 1]),
}
if y_valid.nunique() > 1:
metrics["roc_auc"] = roc_auc_score(y_valid, valid_prob)
fold_metrics.append(metrics)
metric_names = sorted(
metric for metrics in fold_metrics for metric in metrics if metric != "fold"
)
averaged = {}
for metric in metric_names:
values = [fold[metric] for fold in fold_metrics if metric in fold]
averaged[f"cv_mean_{metric}"] = float(np.mean(values))
return averaged
def _save_shap_summary(model, X: pd.DataFrame, path: Path) -> None:
"""Create a SHAP summary plot for the final fitted model."""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import shap
explainer = shap.TreeExplainer(model)
shap_values = _positive_class_shap_values(explainer.shap_values(X))
shap.summary_plot(shap_values, X, show=False)
plt.tight_layout()
plt.savefig(path, dpi=160, bbox_inches="tight")
plt.close()
def train_model(
data_path: str = "data_output/fastf1_races.parquet",
experiment_name: str = "kronector-week3",
n_splits: int = 5,
) -> str:
"""
Train a LightGBM model and log model artifacts to MLflow.
Returns:
MLflow run id.
"""
import mlflow
import mlflow.lightgbm
df = pd.read_parquet(data_path)
bundle, encoders = prepare_model_data(df)
cv_metrics = _cross_validate(bundle, n_splits=n_splits)
final_model = _build_model()
final_model.fit(bundle.X, bundle.y)
tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
if not tracking_uri:
mlflow.set_tracking_uri("file:./mlruns")
mlflow.set_experiment(experiment_name)
with mlflow.start_run() as run:
mlflow.log_params(final_model.get_params())
mlflow.log_params(
{
"n_splits": n_splits,
"n_features": len(bundle.feature_columns),
"n_rows": len(bundle.X),
"model_type": "LightGBM",
}
)
mlflow.log_metrics(cv_metrics)
mlflow.lightgbm.log_model(
final_model,
artifact_path="model",
registered_model_name=REGISTERED_MODEL_NAME,
)
with tempfile.TemporaryDirectory() as tmpdir:
tmp_path = Path(tmpdir)
encoder_path = tmp_path / "label_encoders.pkl"
save_encoders(encoders, str(encoder_path))
mlflow.log_artifact(str(encoder_path), artifact_path="encoders")
shap_path = tmp_path / "shap_summary.png"
_save_shap_summary(final_model, bundle.X, shap_path)
mlflow.log_artifact(str(shap_path), artifact_path="explainability")
return run.info.run_id
def main() -> None:
parser = argparse.ArgumentParser(description="Train KRONECTOR LightGBM model")
parser.add_argument("--data-path", default="data_output/fastf1_races.parquet")
parser.add_argument("--experiment-name", default="kronector-week3")
parser.add_argument("--n-splits", type=int, default=5)
args = parser.parse_args()
run_id = train_model(
data_path=args.data_path,
experiment_name=args.experiment_name,
n_splits=args.n_splits,
)
print(f"MLflow run_id: {run_id}")
if __name__ == "__main__":
main()
|