src/ml/predict.py

"""Inference helper used by the pipeline orchestrator and Gradio app."""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.runtime import configure_runtime  # noqa: E402

configure_runtime()

import joblib
import pandas as pd

from src.config import ML_PIPELINE_PATH  # noqa: E402


class PrepTimePredictor:
    """Load the trained pipeline once and expose ``predict_minutes``."""

    def __init__(self, path: Path | None = None) -> None:
        path = path or ML_PIPELINE_PATH
        if not path.exists():
            raise FileNotFoundError(
                f"Model artifact not found at {path}. Train first with src.ml.train."
            )
        bundle = joblib.load(path)
        self.pipeline = bundle["pipeline"]
        self.feature_cols: list[str] = bundle["feature_cols"]
        self.num_cols: list[str] = bundle["num_cols"]
        self.cat_cols: list[str] = bundle["cat_cols"]
        self.best_model: str = bundle.get("best_model", "unknown")

    def _build_row(self, features: dict[str, Any]) -> pd.DataFrame:
        row = {c: features.get(c) for c in self.feature_cols}
        return pd.DataFrame([row])

    def predict_minutes(self, features: dict[str, Any]) -> float:
        row = self._build_row(features)
        pred = float(self.pipeline.predict(row)[0])
        return max(pred, 1.0)