komitkomit

Dockerfile

@@ -7,4 +7,4 @@ RUN pip install -r requirements.txt
 
 COPY . .
 
-CMD ["python", "app.py"]
+CMD ["gunicorn", "-b", "0.0.0.0:7860", "app:app"]

app.py

@@ -1,10 +1,9 @@
 from flask import Flask, request, jsonify
-import pickle
+import joblib
 
 app = Flask(__name__)
 
 # Load model
-import joblib
 model = joblib.load("models/model.pkl")
 
 @app.route("/")
@@ -13,10 +12,18 @@ def index():
 
 @app.route("/predict", methods=["POST"])
 def predict():
-    data = request.get_json()
-    features = [data["sepal_length"], data["sepal_width"], data["petal_length"], data["petal_width"]]
-    prediction = model.predict([features])
-    return jsonify({"prediction": prediction[0]})
+    try:
+        data = request.get_json()
+        features = [
+            float(data["sepal_length"]),
+            float(data["sepal_width"]),
+            float(data["petal_length"]),
+            float(data["petal_width"]),
+        ]
+        prediction = model.predict([features])
+        return jsonify({"prediction": prediction[0]})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 400
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+# NOTE: Don't run app here if you're using Gunicorn
+# Gunicorn will handle the server initialization