Camilo Piñón
03/05/2024, 7:36 AMfrom typing import Dict, List, Optional
import bentoml
import mlflow
import pandas as pd
MLFLOW_MODEL_NAME = "iforest-test:latest" # TODO: Take this as a parameter
# TODO: Generalize to anomaly detection service
@bentoml.service(resources={"cpu": "200m", "memory": "512Mi"})
class AnomalyDetectionService:
anomaly_detection_model = bentoml.mlflow.get(MLFLOW_MODEL_NAME)
def __init__(self):
mlflow_model_path = self.anomaly_detection_model.path_of(
bentoml.mlflow.MLFLOW_MODEL_FOLDER
)
print(mlflow_model_path)
self.model = mlflow.pyfunc.load_model(mlflow_model_path)
@bentoml.api
def predict(
self,
X: pd.DataFrame,
params: Optional[Dict[str, str]] = {"threshold": "threshold_q95"},
) -> pd.DataFrame:
return self.model.predict(X, params=params)