Update main.py

main.py

@@ -1,42 +1,27 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List
+import onnxruntime as ort
 import numpy as np
-import joblib
-import os
 
-app = FastAPI(title="Stroke Risk Prediction")
+app = FastAPI()
 
-MODEL_PATH = "best_random_forest.pkl"
-_model = None
+# Load the ONNX model session
+session = ort.InferenceSession("app/model.onnx", providers=["CPUExecutionProvider"])
+input_name = session.get_inputs()[0].name
+output_name = session.get_outputs()[0].name
 
-def get_model():
-    global _model
-    if _model is None:
-        if not os.path.exists(MODEL_PATH):
-            raise FileNotFoundError(f"Missing model file: {MODEL_PATH}")
-        _model = joblib.load(MODEL_PATH)
-    return _model
-
-class PredictRequest(BaseModel):
-    input_data: List[List[float]]
-
-@app.get("/")
-def root():
-    return {"status": "ok"}
-
-@app.get("/health")
-def health():
-    return {"status": "healthy"}
+# Define input schema
+class InputData(BaseModel):
+    features: list[float]
 
 @app.post("/predict")
-def predict(req: PredictRequest):
+def predict(data: InputData):
     try:
-        clf = get_model()
-        arr = np.array(req.input_data, dtype=float)
-        preds = clf.predict(arr)
-        return {"predictions": preds.tolist()}
+        input_array = np.array([data.features], dtype=np.float32)
+        prediction = session.run([output_name], {input_name: input_array})
+        return {"prediction": prediction[0][0].item()}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
+