cspj-application/server-ml/main.py
2024-12-02 20:45:50 +08:00

42 lines
1,005 B
Python

from flask import Flask, request, jsonify
import joblib
import os
# first see if the trained model file exists
if os.path.exists("model.pkl"):
model = joblib.load("model.pkl")
else:
raise FileNotFoundError("The model.pkl file does not exist.")
# load the model
model = joblib.load("model.pkl")
# create the flask app
app = Flask(__name__)
# app route for the predict endpoint
@app.route("/predict", methods=["POST"])
def predict():
try:
# get the input query from the request
data = request.json
query = data.get("query", "")
if not query:
return jsonify({"error": "No query provided."}), 400
# make a prediction
# the model expects an array
prediction = model.predict([query])[0]
bad = bool(prediction)
return jsonify({"bad": bad}), 200
except Exception as error:
return jsonify({"error": str(error)}), 500
if __name__ == "__main__":
app.run(debug=True, host="0.0.0.0", port=5000)