Update app.py

app.py

@@ -34,12 +34,10 @@ app.add_middleware(
 # ------------------------------------------------------------
 os.environ['HF_HOME'] = '/tmp/huggingface'
 os.makedirs(os.environ['HF_HOME'], exist_ok=True)
-logger.info(f"HF_HOME set to {os.environ['HF_HOME']}")
 
 MODEL_REPO_ID = "costaspinto/PulmoProbe"
 MODEL_FILENAME = "best_model.joblib"
 
-logger.info("Downloading model from Hugging Face Hub...")
 try:
     model_path = hf_hub_download(
         repo_id=MODEL_REPO_ID,
@@ -47,13 +45,13 @@ try:
         cache_dir=os.environ['HF_HOME']
     )
     model = joblib.load(model_path)
-    logger.info("Model loaded successfully.")
+    logger.info("✅ Model loaded successfully")
 except Exception as e:
-    logger.error(f"Failed to load model: {str(e)}")
+    logger.error(f"❌ Failed to load model: {str(e)}")
     raise RuntimeError(f"Model loading failed: {str(e)}")
 
 # ------------------------------------------------------------
-# Define Input Schema (Corrected Names)
+# Define Input Schema (One-Hot Encoded)
 # ------------------------------------------------------------
 class OneHotPatientData(BaseModel):
     age: float
@@ -65,7 +63,7 @@ class OneHotPatientData(BaseModel):
     other_cancer: int
     gender_Male: int
     family_history_Yes: int
-    
+
     country_Belgium: int
     country_Bulgaria: int
     country_Croatia: int
@@ -93,7 +91,6 @@ class OneHotPatientData(BaseModel):
     country_Spain: int
     country_Sweden: int
 
-    # Corrected to use uppercase Roman numerals
     cancer_stage_Stage_II: int
     cancer_stage_Stage_III: int
     cancer_stage_Stage_IV: int
@@ -101,7 +98,7 @@ class OneHotPatientData(BaseModel):
     smoking_status_Former_Smoker: int
     smoking_status_Never_Smoked: int
     smoking_status_Passive_Smoker: int
-    
+
     treatment_type_Combined: int
     treatment_type_Radiation: int
     treatment_type_Surgery: int
@@ -111,7 +108,7 @@ class OneHotPatientData(BaseModel):
 # ------------------------------------------------------------
 @app.get("/")
 def read_root():
-    return {"message": "Welcome to the PulmoProbe AI API"}
+    return {"message": "Welcome to PulmoProbe AI API"}
 
 # ------------------------------------------------------------
 # Prediction Endpoint
@@ -121,46 +118,38 @@ def predict(data: OneHotPatientData):
     try:
         input_dict = data.dict()
         logger.info(f"Incoming data: {input_dict}")
-        
-        # Define the exact feature order your model expects (with underscores and uppercase Roman numerals)
+
         feature_order = [
             'age', 'bmi', 'cholesterol_level', 'hypertension', 'asthma', 
-            'cirrhosis', 'other_cancer', 'gender_Male', 'country_Belgium', 
-            'country_Bulgaria', 'country_Croatia', 'country_Cyprus', 
-            'country_Czech_Republic', 'country_Denmark', 'country_Estonia', 
-            'country_Finland', 'country_France', 'country_Germany', 
-            'country_Greece', 'country_Hungary', 'country_Ireland', 
-            'country_Italy', 'country_Latvia', 'country_Lithuania', 
-            'country_Luxembourg', 'country_Malta', 'country_Netherlands', 
-            'country_Poland', 'country_Portugal', 'country_Romania', 
-            'country_Slovakia', 'country_Slovenia', 'country_Spain', 
-            'country_Sweden', 
-            'cancer_stage_Stage_II', 'cancer_stage_Stage_III', 'cancer_stage_Stage_IV', 
-            'family_history_Yes', 
-            'smoking_status_Former_Smoker', 'smoking_status_Never_Smoked', 
-            'smoking_status_Passive_Smoker', 'treatment_type_Combined', 
-            'treatment_type_Radiation', 'treatment_type_Surgery'
+            'cirrhosis', 'other_cancer', 'gender_Male', 
+            'country_Belgium','country_Bulgaria','country_Croatia','country_Cyprus',
+            'country_Czech_Republic','country_Denmark','country_Estonia','country_Finland',
+            'country_France','country_Germany','country_Greece','country_Hungary',
+            'country_Ireland','country_Italy','country_Latvia','country_Lithuania',
+            'country_Luxembourg','country_Malta','country_Netherlands','country_Poland',
+            'country_Portugal','country_Romania','country_Slovakia','country_Slovenia',
+            'country_Spain','country_Sweden',
+            'cancer_stage_Stage_II','cancer_stage_Stage_III','cancer_stage_Stage_IV',
+            'family_history_Yes',
+            'smoking_status_Former_Smoker','smoking_status_Never_Smoked','smoking_status_Passive_Smoker',
+            'treatment_type_Combined','treatment_type_Radiation','treatment_type_Surgery'
         ]
 
-        # Create DataFrame and ensure the columns are in the correct order
-        input_df = pd.DataFrame([input_dict], columns=feature_order)
-        logger.info(f"DataFrame for prediction: {input_df}")
+        # Fill missing fields with 0
+        input_dict_complete = {col: input_dict.get(col, 0) for col in feature_order}
+        input_df = pd.DataFrame([input_dict_complete], columns=feature_order)
 
         # Predict probabilities
         probabilities = model.predict_proba(input_df)[0]
-        logger.info(f"Model probabilities: {probabilities}")
-
-        confidence_high_risk = probabilities[0]
+        confidence_high_risk = probabilities[1]  # Class 1 = High Risk
         risk_level = "High Risk of Non-Survival" if confidence_high_risk > 0.5 else "Low Risk of Non-Survival"
 
         result = {
             "risk": risk_level,
-            "confidence": f"{confidence_high_risk * 100:.1f}%"
+            "confidence": f"{confidence_high_risk*100:.1f}%"
         }
-
-        logger.info(f"Prediction result: {result}")
         return result
 
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")
-        return {"error": str(e), "input_data_received": data.dict()}
+        return {"error": str(e), "input_data_received": data.dict()}