blob: 410f7cceb2fa2791a28d063d47571bab2b43b7c5 [file] [log] [blame]
#!/usr/bin/env python
# coding=utf-8
"""Predictor engine action.
Use this module to add the project main code.
"""
from .._compatibility import six
from .._logging import get_logger
from marvin_python_toolbox.engine_base import EngineBasePrediction
__all__ = ['Predictor']
logger = get_logger('predictor')
class Predictor(EngineBasePrediction):
def __init__(self, **kwargs):
super(Predictor, self).__init__(**kwargs)
def execute(self, input_message, params, **kwargs):
final_prediction = {
"prediction_rf": self.marvin_model['rf'].predict([input_message])[0],
"prediction_svm": self.marvin_model['svm'].predict([input_message])[0]
}
print(final_prediction)
return final_prediction