import pickle
import os


class PredictionModel:
    model = None

    def __init__(self, id_model="rf_model"):

        with open(os.path.join(os.path.dirname(__file__), id_model + ".pkl"), 'rb') as fin:
            self.model = pickle.load(fin)

    def predict(self, testing_data):
        prediction = self.model.predict(testing_data)
        return prediction

    def convert_to_json(self, df):
        return df.to_json()