churn-model / handler.py
rajistics's picture
Update handler.py (#1)
1e43c5c
import pandas as pd
import pickle
from typing import Dict, List, Any
import numpy as np
import os
# set device
class EndpointHandler():
def __init__(self, path=""):
# load the optimized model
pathb = os.path.join(path,"./churn.pkl")
self.pipe = pd.read_pickle(pathb)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`list`:. A string representing what the label/class is
"""
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
df = pd.DataFrame(inputs)
df["TotalCharges"] = df["TotalCharges"].replace(" ", np.nan, regex=False).astype(float)
df = df.drop(columns=["customerID"])
df = df.drop(columns=["Churn"])
# run inference pipeline
pred = self.pipe.predict(df)
# postprocess the prediction
return {"pred": pred}