handler.py

import pickle
import numpy as np
from flask import Flask, request, jsonify

# Load the pickle model
MODEL_PATH = "/mnt/data/model.pkl"

app = Flask(__name__)

# Load the model
with open(MODEL_PATH, 'rb') as file:
    model = pickle.load(file)

@app.route('/predict', methods=['POST'])
def predict():
    try:
        # Parse input JSON data
        input_data = request.get_json()
        if not input_data:
            return jsonify({"error": "Invalid input data"}), 400
        
        # Assuming the input data is a list of features
        features = np.array(input_data['features']).reshape(1, -1)  # Adjust for single input

        # Make predictions
        prediction = model.predict(features)
        
        # Return the prediction as JSON
        return jsonify({"prediction": prediction.tolist()}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

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