Ranking Model
import pickle
from pathlib import Path
import pandas as pd
PACKAGE_PATH = Path(__file__).parent
class ModelPackage:
def __init__(self):
self.output_columns = ['score']
with open(f'{PACKAGE_PATH}/model.pkl', 'rb') as infile:
self.model = pickle.load(infile)
def predict(self, input_df):
pred = self.model.predict(input_df)
return pd.DataFrame(pred, columns=self.output_columns)
def get_model():
return ModelPackage()Last updated
Was this helpful?