Upload app.py with huggingface_hub

app.py

@@ -57,11 +57,11 @@ class BatchPredictionRequest(BaseModel):
 class DayPrediction(BaseModel):
     day: int
     prediction: int
-    probability: float
+    probability: float  # Probability of HEADACHE (class 1), regardless of prediction
 
 class SinglePredictionResponse(BaseModel):
     prediction: int
-    probability: float
+    probability: float  # Probability of HEADACHE (class 1), regardless of prediction
 
 class BatchPredictionResponse(BaseModel):
     predictions: List[DayPrediction]
@@ -105,13 +105,22 @@ def predict(request: SinglePredictionRequest):
         # Convert input to numpy array
         features = np.array(request.features).reshape(1, -1)
 
-        # Make prediction
-        prediction = model.predict(features)[0]
-        probability = float(model.predict_proba(features)[0][int(prediction)])
+        # Get probability array for both classes
+        prob_array = model.predict_proba(features)[0]
+
+        # Always return probability of headache (class 1)
+        headache_probability = float(prob_array[1])
+
+        # Make prediction using threshold if available
+        if isinstance(model, dict) and 'optimal_threshold' in model:
+            threshold = model['optimal_threshold']
+            prediction = 1 if headache_probability >= threshold else 0
+        else:
+            prediction = model.predict(features)[0]
 
         return SinglePredictionResponse(
             prediction=int(prediction),
-            probability=probability
+            probability=headache_probability
         )
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Prediction error: {str(e)}")
@@ -129,17 +138,26 @@ def predict_batch(request: BatchPredictionRequest):
         if features.ndim != 2:
             raise ValueError(f"Expected 2D array, got shape {features.shape}")
 
-        # Make predictions for all days
-        predictions = model.predict(features)
+        # Get probabilities for all days
         probabilities = model.predict_proba(features)
 
         # Format results
         results = []
-        for i, (pred, prob_array) in enumerate(zip(predictions, probabilities), 1):
+        for i, prob_array in enumerate(probabilities, 1):
+            # Always use probability of headache (class 1)
+            headache_probability = float(prob_array[1])
+
+            # Make prediction using threshold if available
+            if isinstance(model, dict) and 'optimal_threshold' in model:
+                threshold = model['optimal_threshold']
+                prediction = 1 if headache_probability >= threshold else 0
+            else:
+                prediction = model.predict(features[i-1:i])[0]
+
             results.append(DayPrediction(
                 day=i,
-                prediction=int(pred),
-                probability=float(prob_array[int(pred)])
+                prediction=int(prediction),
+                probability=headache_probability
             ))
 
         return BatchPredictionResponse(predictions=results)