Spaces:
Sleeping
Sleeping
File size: 4,320 Bytes
d623240 | 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 | from __future__ import annotations
import logging
import sys
from dataclasses import dataclass
import json
from pathlib import Path
from threading import Lock
from typing import Any
import joblib
import numpy as np
import pandas as pd
import app.ml.custom_transformers as custom_transformers
from app.core.exceptions import ModelInferenceError, ServiceUnavailableError
logger = logging.getLogger("oraculo_api.model")
@dataclass(slots=True)
class ModelPrediction:
label: str
probability: float
raw_probabilities: list[float]
model_version: str
class ModelManager:
def __init__(self, model_path: str | Path):
self.model_path = Path(model_path)
self.pipeline: Any = None
self.manifest: dict[str, Any] = {}
self._lock = Lock()
@property
def is_loaded(self) -> bool:
return self.pipeline is not None
@property
def model_version(self) -> str:
if self.manifest.get("model_version"):
return str(self.manifest["model_version"])
if self.pipeline is None:
return "unloaded"
return str(getattr(self.pipeline, "version", "unknown"))
@property
def manifest_path(self) -> Path:
return self.model_path.with_name("model_manifest.json")
def _register_pickle_bridge(self) -> None:
setattr(
sys.modules["__main__"],
"PipelineProduccionMLOps",
custom_transformers.PipelineProduccionMLOps,
)
def load_model(self) -> None:
with self._lock:
if self.pipeline is not None:
return
if not self.model_path.exists():
raise ServiceUnavailableError(f"Model artifact not found at '{self.model_path}'.")
try:
self._register_pickle_bridge()
self.pipeline = joblib.load(self.model_path)
if hasattr(self.pipeline, "_infer_missing_artefacts"):
self.pipeline._infer_missing_artefacts()
if self.manifest_path.exists():
self.manifest = json.loads(self.manifest_path.read_text(encoding="utf-8"))
logger.info("Model artifact loaded from %s", self.model_path)
except Exception as exc:
logger.exception("Unable to load model artifact: %s", exc)
raise ServiceUnavailableError("Model artifact could not be loaded.") from exc
def unload_model(self) -> None:
with self._lock:
self.pipeline = None
self.manifest = {}
def _ensure_loaded(self) -> Any:
if self.pipeline is None:
raise ServiceUnavailableError("Model is not loaded.")
return self.pipeline
@staticmethod
def _to_frame(input_data: dict[str, Any] | pd.DataFrame) -> pd.DataFrame:
if isinstance(input_data, pd.DataFrame):
return input_data.copy()
return pd.DataFrame([input_data])
def predict(self, input_data: dict[str, Any] | pd.DataFrame) -> np.ndarray:
pipeline = self._ensure_loaded()
try:
frame = self._to_frame(input_data)
return pipeline.predict(frame)
except Exception as exc:
logger.exception("Prediction failed: %s", exc)
raise ModelInferenceError("Prediction failed.") from exc
def predict_proba(self, input_data: dict[str, Any] | pd.DataFrame) -> np.ndarray:
pipeline = self._ensure_loaded()
try:
frame = self._to_frame(input_data)
return pipeline.predict_proba(frame)
except Exception as exc:
logger.exception("Probability inference failed: %s", exc)
raise ModelInferenceError("Probability inference failed.") from exc
def predict_one(self, input_data: dict[str, Any]) -> ModelPrediction:
labels = self.predict(input_data)
probabilities = self.predict_proba(input_data)
probability_vector = probabilities[0].tolist()
positive_probability = float(probability_vector[1] if len(probability_vector) > 1 else probability_vector[0])
return ModelPrediction(
label=str(labels[0]),
probability=positive_probability,
raw_probabilities=probability_vector,
model_version=self.model_version,
)
|