Update app.py

app.py

@@ -3,8 +3,9 @@ from pydantic import BaseModel
 import joblib
 import numpy as np
 from fastapi.middleware.cors import CORSMiddleware
+import pandas as pd # Needed if you decide to preprocess raw input in the API
 
-app = FastAPI(title="Web Attack Detection API")
+app = FastAPI(title="Multi-Class Web Data Classifier API")
 
 # Allow browser requests
 app.add_middleware(
@@ -15,22 +16,28 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Expected number of features
-EXPECTED_FEATURES = 35
+# The number of features MUST match the X_preprocessed.shape[1] from training
+# Our X_preprocessed had 5 columns: 'feature1', 'id', 'type_defacement', 'type_malware', 'type_phishing'
+EXPECTED_FEATURES = 5
 
-# Request schema
+# Request schema for the preprocessed features
+# IMPORTANT: The client calling this API must preprocess their data
+# (impute, scale numerical, one-hot encode categorical) to match the format 
+# of X_preprocessed used during model training.
 class InputData(BaseModel):
-    features: list[float]
+    features: list[float] # Expecting 5 preprocessed float values
 
-# Load model
+# Load model and label encoder
 try:
-    model = joblib.load("web_attack_detection_model.pkl")
+    model = joblib.load("trained_lightgbm_model.joblib")
+    label_encoder = joblib.load("label_encoder.pkl")
+    print("Model and Label Encoder loaded successfully.")
 except Exception as e:
-    raise RuntimeError(f"Model failed to load: {e}")
+    raise RuntimeError(f"Error loading model or label encoder: {e}")
 
 @app.get("/")
 def home():
-    return {"message": "Web attack detection model running"}
+    return {"message": "Multi-Class Web Data Classifier API running"}
 
 @app.get("/health")
 def health():
@@ -38,24 +45,23 @@ def health():
 
 @app.post("/predict")
 def predict(data: InputData):
-
     if len(data.features) != EXPECTED_FEATURES:
         raise HTTPException(
             status_code=400,
-            detail=f"Expected {EXPECTED_FEATURES} features"
+            detail=f"Expected {EXPECTED_FEATURES} features, but received {len(data.features)}"
         )
 
-    input_data = np.array(data.features).reshape(1, -1)
+    # Convert input list to numpy array and reshape for prediction
+    # This assumes input features are ALREADY preprocessed as per training
+    input_array = np.array(data.features).reshape(1, -1)
 
-    prediction = model.predict(input_data)[0]
+    # Make prediction
+    prediction_encoded = model.predict(input_array)[0]
 
-    # Convert prediction to readable output
-    if prediction == -1:
-        result = "attack detected"
-    else:
-        result = "normal request"
+    # Inverse transform to get the original label string
+    prediction_label = label_encoder.inverse_transform([prediction_encoded])[0]
 
     return {
-        "prediction": result,
-        "raw_prediction": int(prediction)
-    }
+        "predicted_label": prediction_label,
+        "raw_prediction_encoded": int(prediction_encoded)
+    }