Update api.py

api.py

@@ -1,3 +1,4 @@
+from pydantic import BaseModel
 from huggingface_hub import hf_hub_download
 import joblib
 from fastapi import FastAPI
@@ -7,23 +8,27 @@ import os
 # Set a different cache directory
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 
-# Initialize FastAPI app
 app = FastAPI()
 
-# Download the model from Hugging Face Model Repo
+# Download model
 model_path = hf_hub_download(
-    repo_id="haseebnawazz/sleep-stage-classifier-model",   # Update with your HF model repo
+    repo_id="haseebnawazz/sleep-stage-classifier-model",
     filename="model.joblib",
-    cache_dir="/tmp/hf_cache"  # Ensure the cache is stored in a writable directory
+    cache_dir="/tmp/hf_cache"
 )
 model = joblib.load(model_path)
 
-# Define a prediction endpoint
+# ✅ Define Pydantic schema to accept {"features": {...}}
+class FeaturesInput(BaseModel):
+    features: dict
+
+# ✅ Update endpoint to use schema
 @app.post("/predict")
-def predict(features: list):
-    features = np.array(features).reshape(1, -1)  # Convert input to NumPy array
-    prediction = model.predict(features).tolist()  # Predict and convert to list
-    return {"prediction": prediction}
+def predict(input: FeaturesInput):
+    features = input.features
+    features_array = np.array([list(features.values())])  # Ensure 2D array for model
+    prediction = model.predict(features_array).tolist()
+    return {"predicted_stage": prediction[0]}
 
 @app.get("/")
 async def read_root():
@@ -31,4 +36,4 @@ async def read_root():
 
 @app.get("/logs")
 async def get_logs():
-    return {"logs": "container logs here"}
+    return {"logs": "container logs here"}