Update main.py

main.py

@@ -1,56 +1,71 @@
 import joblib
 from pydantic import BaseModel
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 import uvicorn
 import logging
+import os
+
+logging.basicConfig(level=logging.INFO)
 
-logging.basicConfig(level = logging.INFO)
 # 1. Load the trained model
-model = joblib.load('frauddetection.pkl')
+try:
+    if os.path.exists('frauddetection.pkl'):
+        model = joblib.load('frauddetection.pkl')
+    else:
+        raise FileNotFoundError("Model file 'frauddetection.pkl' not found.")
+except Exception as e:
+    logging.error(f"Error loading model: {e}")
+    model = None
 
 # 2. Define the input data schema using Pydantic BaseModel
 class InputData(BaseModel):
-    Year:int
-    Month:int
-    UseChip:int
-    Amount:int
-    MerchantName:int
-    MerchantCity:int
-    MerchantState:int
-    mcc:int
-    # Add the rest of the input features (feature4, feature5, ..., feature12)
+    Year: int
+    Month: int
+    UseChip: int
+    Amount: int
+    MerchantName: int
+    MerchantCity: int
+    MerchantState: int
+    mcc: int
 
 # 3. Create a FastAPI app
 app = FastAPI()
+
 @app.get('/')
 def welcome():
     return {"Welcome": "This is the home page of the API"}
 
-
 # 4. Define the prediction route
 @app.post('/predict/')
 async def predict(data: InputData):
+    if model is None:
+        raise HTTPException(status_code=500, detail="Model not loaded properly.")
+
     # Convert the input data to a dictionary
     input_data = data.dict()
 
     # Extract the input features from the dictionary
-    feature1 = input_data['Year']
-    feature2=input_data['Month']
-    feature3=input_data['UseChip']
-    feature4=input_data['Amount']
-    feature5=input_data['MerchantName']
-    feature6=input_data['MerchantCity']
-    feature7=input_data['MerchantState']
-    feature8=input_data['mcc']
-    # Extract the rest of the input features (feature4, feature5, ..., feature12)
-
-    # Perform the prediction using the loaded model
-    prediction = model.predict([[feature1, feature2, feature3,feature4,feature5,feature6,feature7,feature8]])  # Replace ... with the rest of the features
-    # Convert the prediction to a string (or any other format you prefer)
-    result = "Fraud" if prediction[0] == 1 else "Not a Fraud"
-    # result = prediction 
-    return {"prediction": result}
-# 4. Run the API with uvicorn
-#    Will run on http://127.0.0.1:8000
+    feature_list = [
+        input_data['Year'],
+        input_data['Month'],
+        input_data['UseChip'],
+        input_data['Amount'],
+        input_data['MerchantName'],
+        input_data['MerchantCity'],
+        input_data['MerchantState'],
+        input_data['mcc']
+    ]
+
+    try:
+        # Perform the prediction using the loaded model
+        prediction = model.predict([feature_list])  # Ensure model expects the correct number of features
+        result = "Fraud" if prediction[0] == 1 else "Not a Fraud"
+        return {"prediction": result}
+    except Exception as e:
+        logging.error(f"Prediction error: {e}")
+        raise HTTPException(status_code=500, detail="An error occurred during prediction.")
+
+# 5. Run the API with uvicorn
+#    Will run on http://127.0.0.1:8080
 if __name__ == '__main__':
-    uvicorn.run(app, port=8080)
+    uvicorn.run(app, host="127.0.0.1", port=8080)