File size: 3,066 Bytes
2532605 a13ee72 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 | """
KRONECTOR - Prediction helpers.
Inference must reuse the LabelEncoders fitted during training. Do not call
prepare_model_data without the loaded encoders in production prediction paths.
"""
from __future__ import annotations
import argparse
import numpy as np
import pandas as pd
from ml.feature_engineering import load_encoders, prepare_model_data
def load_model_and_encoders(run_id: str):
"""Load a logged MLflow model and its fitted categorical encoders."""
import mlflow.lightgbm
from mlflow.artifacts import download_artifacts
import logging
logger = logging.getLogger(__name__)
try:
model = mlflow.lightgbm.load_model(f"runs:/{run_id}/model")
except Exception as e:
logger.warning(f"Could not load from runs:/ URI: {e}. Trying registered model...")
try:
model = mlflow.lightgbm.load_model("models:/kronector-f1-lgbm/latest")
except Exception as e2:
logger.error(f"Failed to load registered model: {e2}")
raise
encoder_path = download_artifacts(
run_id=run_id, artifact_path="encoders/label_encoders.pkl"
)
encoders = load_encoders(encoder_path)
return model, encoders
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 _explain_predictions(model, X: pd.DataFrame) -> list[dict[str, float]]:
import shap
explainer = shap.TreeExplainer(model)
shap_values = _positive_class_shap_values(explainer.shap_values(X))
return [
{
feature_name: float(shap_value)
for feature_name, shap_value in zip(X.columns, row_values)
}
for row_values in shap_values
]
def predict_dataframe(
df: pd.DataFrame, model, encoders: dict, explain: bool = True
) -> pd.DataFrame:
"""Return win probabilities and optional SHAP dictionaries."""
bundle, _ = prepare_model_data(df, encoders=encoders)
if hasattr(model, "predict_proba"):
probabilities = model.predict_proba(bundle.X)[:, 1]
else:
probabilities = model.predict(bundle.X)
result = bundle.metadata.copy()
result["win_probability"] = probabilities
if explain:
result["shap_values"] = _explain_predictions(model, bundle.X)
return result
def main() -> None:
parser = argparse.ArgumentParser(description="Run KRONECTOR model inference")
parser.add_argument("--run-id", required=True)
parser.add_argument("--data-path", default="data_output/fastf1_races.parquet")
args = parser.parse_args()
model, encoders = load_model_and_encoders(args.run_id)
df = pd.read_parquet(args.data_path)
predictions = predict_dataframe(df, model, encoders)
print(predictions.head().to_string(index=False))
if __name__ == "__main__":
main()
|