56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import joblib
|
|
import pandas as pd
|
|
|
|
from app.ml.features import build_feature_matrix
|
|
|
|
|
|
DEFAULT_MODEL_PATH = Path("models/scalping_model.joblib")
|
|
|
|
|
|
class ScalpingModel:
|
|
def __init__(self, model_path: str | Path = DEFAULT_MODEL_PATH):
|
|
self.model_path = Path(model_path)
|
|
self.bundle: dict[str, Any] | None = None
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
return self.model_path.exists()
|
|
|
|
def load(self) -> bool:
|
|
if not self.available:
|
|
return False
|
|
self.bundle = joblib.load(self.model_path)
|
|
return True
|
|
|
|
@property
|
|
def version(self) -> str:
|
|
if self.bundle is None and not self.load():
|
|
return ""
|
|
assert self.bundle is not None
|
|
return str(self.bundle.get("created_at", ""))
|
|
|
|
def score(self, row: dict[str, Any]) -> dict[str, float]:
|
|
if self.bundle is None and not self.load():
|
|
return {}
|
|
|
|
assert self.bundle is not None
|
|
frame = pd.DataFrame([row])
|
|
features, _ = build_feature_matrix(
|
|
frame,
|
|
self.bundle["feature_columns"],
|
|
self.bundle.get("medians"),
|
|
)
|
|
|
|
scores: dict[str, float] = {}
|
|
for target, model in self.bundle.get("models", {}).items():
|
|
if hasattr(model, "predict_proba"):
|
|
scores[target] = float(model.predict_proba(features)[0][1])
|
|
else:
|
|
scores[target] = float(model.predict(features)[0])
|
|
return scores
|