inference/predict.py

import json
from pathlib import Path

import pandas as pd
import xgboost as xgb


class Predictor:

    def __init__(
        self,
        model_path: Path,
        metadata_path: Path,
    ):

        self.model = xgb.XGBClassifier()

        self.model.load_model(
            model_path.as_posix()
        )

        with open(metadata_path) as f:
            self.metadata = json.load(f)

        self.feature_cols = self.metadata[
            "feature_cols"
        ]

    def predict(
        self,
        df: pd.DataFrame,
    ) -> pd.DataFrame:

        missing = [
            c
            for c in self.feature_cols
            if c not in df.columns
        ]

        if missing:

            raise RuntimeError(
                f"Missing model features "
                f"({len(missing)}): {missing}"
            )

        X = df[
            self.feature_cols
        ].copy()

        # Critical contract check
        if list(X.columns) != self.feature_cols:

            raise RuntimeError(
                "Feature order does not match "
                "model_metadata.json"
            )

        probabilities = self.model.predict_proba(
            X
        )[:, 1]

        result = df[
            ["timestamp", "symbol", "close"]
        ].copy()

        result["predicted_probability"] = (
            probabilities
        )

        return result