Update app.py

app.py

@@ -1,13 +1,31 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from __future__ import annotations
+
+import os
+import sys
+import logging
+from pathlib import Path
+from typing import List
+
 import joblib
 import numpy as np
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-import pandas as pd # You might need this if you decide to preprocess raw input in the API
+from pydantic import BaseModel, Field
 
-app = FastAPI(title="Multi-Class Web Data Classifier API")
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='[%(asctime)s] %(name)s — %(levelname)s: %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="CyHub Model 4 — Domain Classification",
+    description="Multi-class web data classifier (Normal/Adult/Betting/Malware)",
+    version="1.0.0"
+)
 
-# Allow browser requests
+# CORS configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,52 +34,237 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 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'
+# ─────────────────────────────────────────────────────────────────
+# Configuration
+# ─────────────────────────────────────────────────────────────────
+
+# Expected feature count (must match training data)
 EXPECTED_FEATURES = 5
 
-# 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.
+# Model file paths (support both relative and absolute)
+MODEL_DIR = Path(os.getenv("MODEL_DIR", "."))
+MODEL_PATH = MODEL_DIR / "trained_lightgbm_model.pkl"
+ENCODER_PATH = MODEL_DIR / "label_encoder.pkl"
+
+logger.info(f"Looking for model at: {MODEL_PATH.absolute()}")
+logger.info(f"Looking for encoder at: {ENCODER_PATH.absolute()}")
+
+# ─────────────────────────────────────────────────────────────────
+# Request/Response Models
+# ─────────────────────────────────────────────────────────────────
+
 class InputData(BaseModel):
-    features: list[float] # Expecting 5 preprocessed float values
+    """Input data for model prediction.
+
+    Must be preprocessed (scaled, imputed, encoded) to match training data.
+    """
+    features: List[float] = Field(
+        ...,
+        min_items=EXPECTED_FEATURES,
+        max_items=EXPECTED_FEATURES,
+        description=f"Exactly {EXPECTED_FEATURES} preprocessed float values"
+    )
+
+
+class PredictionResponse(BaseModel):
+    """Model prediction response."""
+    predicted_label: str
+    raw_prediction_encoded: int
+
+
+# ─────────────────────────────────────────────────────────────────
+# Model Loading
+# ─────────────────────────────────────────────────────────────────
+
+model = None
+label_encoder = None
+model_loaded = False
+
+def load_models():
+    """Load LightGBM model and label encoder."""
+    global model, label_encoder, model_loaded
+
+    try:
+        # Check if files exist
+        if not MODEL_PATH.exists():
+            raise FileNotFoundError(f"Model file not found: {MODEL_PATH.absolute()}")
+        if not ENCODER_PATH.exists():
+            raise FileNotFoundError(f"Encoder file not found: {ENCODER_PATH.absolute()}")
+
+        # Load model
+        logger.info(f"Loading model from {MODEL_PATH.absolute()}...")
+        model = joblib.load(str(MODEL_PATH))
+        logger.info("✓ Model loaded successfully")
+
+        # Load label encoder
+        logger.info(f"Loading label encoder from {ENCODER_PATH.absolute()}...")
+        label_encoder = joblib.load(str(ENCODER_PATH))
+        logger.info("✓ Label encoder loaded successfully")
+
+        # Verify encoder has classes
+        if not hasattr(label_encoder, 'classes_'):
+            raise ValueError("Label encoder missing 'classes_' attribute")
+
+        logger.info(f"Label classes: {label_encoder.classes_.tolist()}")
+        model_loaded = True
+
+        return True
+
+    except FileNotFoundError as e:
+        logger.error(f"File error: {e}")
+        logger.warning(f"Make sure model files are in: {MODEL_DIR.absolute()}")
+        return False
+    except Exception as e:
+        logger.error(f"Error loading models: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+# ─────────────────────────────────────────────────────────────────
+# Startup/Shutdown Events
+# ─────────────────────────────────────────────────────────────────
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize models on startup."""
+    global model_loaded
+    logger.info("Starting up CyHub Model 4 API...")
+
+    if not load_models():
+        logger.error("Failed to load models. API will return 503 until models are available.")
+        model_loaded = False
+    else:
+        logger.info("Model 4 API ready!")
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """Cleanup on shutdown."""
+    logger.info("Shutting down CyHub Model 4 API...")
 
-# Load model and label encoder
-try:
-    model = joblib.load("trained_lightgbm_model.pkl")
-    label_encoder = joblib.load("label_encoder.pkl")
-    print("Model and Label Encoder loaded successfully.")
-except Exception as e:
-    raise RuntimeError(f"Error loading model or label encoder: {e}")
+
+# ─────────────────────────────────────────────────────────────────
+# API Endpoints
+# ─────────────────────────────────────────────────────────────────
 
 @app.get("/")
-def home():
-    return {"message": "Multi-Class Web Data Classifier API running"}
+async def root():
+    """Root endpoint — health check."""
+    return {
+        "service": "CyHub Model 4 — Domain Classification",
+        "status": "healthy",
+        "model_loaded": model_loaded,
+        "version": "1.0.0"
+    }
+
 
 @app.get("/health")
-def health():
-    return {"status": "ok"}
+async def health_check():
+    """Health check endpoint."""
+    if not model_loaded:
+        raise HTTPException(
+            status_code=503,
+            detail="Model not loaded. Check server logs."
+        )
+
+    return {
+        "status": "healthy",
+        "model_status": "ready",
+        "expected_features": EXPECTED_FEATURES
+    }
+
+
+@app.post("/predict", response_model=PredictionResponse)
+async def predict(data: InputData) -> PredictionResponse:
+    """Predict domain classification from preprocessed features.
+
+    Args:
+        data: InputData with exactly 5 preprocessed features
 
-@app.post("/predict")
-def predict(data: InputData):
-    if len(data.features) != EXPECTED_FEATURES:
+    Returns:
+        PredictionResponse with predicted_label and raw_prediction_encoded
+
+    Raises:
+        HTTPException: If model not loaded or feature count incorrect
+    """
+    # Check if model is loaded
+    if not model_loaded or model is None or label_encoder is None:
+        logger.error("Model not loaded, cannot make prediction")
+        raise HTTPException(
+            status_code=503,
+            detail="Model not loaded. Check server logs."
+        )
+
+    try:
+        # Validate feature count
+        if len(data.features) != EXPECTED_FEATURES:
+            raise ValueError(
+                f"Expected {EXPECTED_FEATURES} features, got {len(data.features)}"
+            )
+
+        # Validate all features are float/convertible
+        features_array = np.array(data.features, dtype=float)
+
+        # Check for NaN or Inf
+        if not np.all(np.isfinite(features_array)):
+            raise ValueError("Features contain NaN or Inf values")
+
+        # Reshape for prediction (1 sample, 5 features)
+        input_array = features_array.reshape(1, -1)
+
+        logger.debug(f"Input features: {data.features}")
+
+        # Make prediction
+        prediction_encoded = model.predict(input_array)[0]
+        logger.debug(f"Raw prediction: {prediction_encoded}")
+
+        # Ensure prediction is valid
+        if prediction_encoded not in label_encoder.classes_:
+            logger.warning(f"Unexpected prediction value: {prediction_encoded}")
+
+        # Decode prediction to label
+        try:
+            prediction_label = label_encoder.inverse_transform([prediction_encoded])[0]
+        except Exception as e:
+            logger.error(f"Error decoding prediction: {e}")
+            raise ValueError(f"Could not decode prediction {prediction_encoded}")
+
+        logger.info(f"Prediction: {prediction_label}")
+
+        return PredictionResponse(
+            predicted_label=prediction_label,
+            raw_prediction_encoded=int(prediction_encoded)
+        )
+
+    except ValueError as e:
+        logger.error(f"Validation error: {e}")
         raise HTTPException(
             status_code=400,
-            detail=f"Expected {EXPECTED_FEATURES} features, but received {len(data.features)}"
+            detail=str(e)
+        )
+    except Exception as e:
+        logger.error(f"Prediction error: {e}")
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(
+            status_code=500,
+            detail=f"Prediction failed: {str(e)}"
         )
 
-    # 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)
 
-    # Make prediction
-    prediction_encoded = model.predict(input_array)[0]
+# ─────────────────────────────────────────────────────────────────
+# Development Entry Point
+# ─────────────────────────────────────────────────────────────────
 
-    # Inverse transform to get the original label string
-    prediction_label = label_encoder.inverse_transform([prediction_encoded])[0]
+if __name__ == "__main__":
+    import uvicorn
 
-    return {
-        "predicted_label": prediction_label,
-        "raw_prediction_encoded": int(prediction_encoded)
-    }
+    # For development/testing
+    logger.info("Starting Model 4 API in development mode...")
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,  # Default HuggingFace Spaces port
+        log_level="info"
+    )