csic_model / handler.py
Anshika2002's picture
Upload handler.py
c588c04 verified
import joblib
import pandas as pd
class EndpointHandler:
def __init__(self, path=""):
# load your model artifact
self.artifact = joblib.load(f"{path}/csic_rf_model.joblib")
def __call__(self, data):
# input must be a list in correct feature order
inputs = data["inputs"]
df = pd.DataFrame([inputs], columns=self.artifact["features"])
pred = self.artifact["model"].predict(df)
# if target encoder exists
if "target_encoder" in self.artifact:
label = self.artifact["target_encoder"].inverse_transform(pred)
return {"prediction": label[0]}
return {"prediction": int(pred[0])}