from flask import Flask, jsonify, request
import mlflow.pyfunc

app = Flask(__name__)

# Load the model as a PyFuncModel.
model = mlflow.pyfunc.load_model(model_uri="models:/deployed_model/1")

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json()
    predictions = model.predict(data)
    return jsonify(predictions.tolist())

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)