Update main.py

main.py

@@ -53,12 +53,13 @@ async def root():
 
 @app.post("/predict/")
 async def predict_income(data: IncomePredictionRequest):
-    input_data = data.dict()
-    input_df = pd.DataFrame([input_data])
-    prediction = dt_model.predict(input_df)
-    prediction_proba = dt_model.predict_proba(input_df)
-    prediction_result = "Income over $50K" if prediction[0] == 1 else "Income under $50K"
-    return {"income_prediction": prediction_result, "prediction_probability": prediction_proba[0][1]}
+    try:
+        input_data = data.dict()
+        input_df = pd.DataFrame([input_data])
+        prediction = dt_model.predict(input_df)
+        prediction_proba = dt_model.predict_proba(input_df)
+        prediction_result = "Income over $50K" if prediction[0] == 1 else "Income under $50K"
+        return {"income_prediction": prediction_result, "prediction_probability": prediction_proba[0][1]}
 
     except Exception as e:
         logging.error(f"Prediction failed: {e}")