Update inference.py

inference.py

@@ -1,18 +1,34 @@
 import pickle
 import pandas as pd
+import json
 
 def load_model():
-    with open("model/expense_forecaster_model.pkl", "rb") as f:
-        model = pickle.load(f)
-    return model
+    try:
+        with open("model/expense_forecaster_model.pkl", "rb") as f:
+            model = pickle.load(f)
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
 
 def predict(data):
     model = load_model()
-    df = pd.DataFrame([data])
-    prediction = model.predict(df)
-    return prediction.tolist()
+    if model is None:
+        return {"error": "Model loading failed"}
+
+    try:
+        # Ensure data is a dictionary
+        if not isinstance(data, dict):
+            return {"error": "Input data must be a dictionary"}
+
+        df = pd.DataFrame([data])
+        prediction = model.predict(df)
+        return prediction.tolist()
+
+    except Exception as e:
+        return {"error": f"Prediction error: {e}"}
 
 if __name__ == "__main__":
-    example_input = {"income": 5000, "previous_expenses": 3000, "month": 12} #example data, change this to match your feature names.
+    example_input = {"income": 5000, "previous_expenses": 3000, "month": 12}
     prediction = predict(example_input)
     print(f"Prediction: {prediction}")