Update app.py

app.py

@@ -3,119 +3,104 @@ import torch
 import numpy as np
 
 # -----------------------------
-# 1. Load Your Specific Model
+# 1. Load Model (Robust)
 # -----------------------------
 MODEL_PATH = "models/phishing_rf_model.pt"
 
-print(f"Loading model from {MODEL_PATH}...")
-
-# We use torch.load because the file extension is .pt
-# map_location='cpu' ensures it works on servers without massive GPUs
+print(f"Attempting to load model from {MODEL_PATH}...")
 try:
+    # Load the model file
     model = torch.load(MODEL_PATH, map_location=torch.device('cpu'))
-    print("✅ Model loaded successfully!")
+    print(f"✅ Model loaded successfully! Type: {type(model)}")
 except Exception as e:
     print(f"❌ Failed to load model: {e}")
-    raise e
-
-# -----------------------------
-# 2. Feature Extraction (No Scaler)
-# -----------------------------
-def extract_features(url: str) -> np.ndarray:
-    """
-    Extracts the features the model expects.
-    Since we are skipping the scaler, we feed these raw numbers directly.
-    """
-    length = len(url)
-    dots = url.count('.')
-    hyphens = url.count('-')
-    digits = sum(c.isdigit() for c in url)
-    at_sign = url.count('@')
-    
-    # Create the array shape [1, 5] (1 sample, 5 features)
-    return np.array([[length, dots, hyphens, digits, at_sign]], dtype=float)
+    model = None
 
 # -----------------------------
-# 3. Prediction Logic
+# 2. Prediction Logic (Universal)
 # -----------------------------
-def predict_phishing(url: str):
+def predict_phishing(url):
+    # Safety checks
+    if model is None:
+        return {"Error": 0.0}, "Model failed to load. Check logs."
     if not url:
-        return None
+        return None, "Please enter a URL."
 
-    # 1. Extract features
-    features = extract_features(url)
-    
-    # 2. Predict
-    # We assume the model inside the .pt file is a standard sklearn model 
-    # (RandomForest) that supports .predict_proba()
     try:
-        pred_prob = model.predict_proba(features)[0]
-    except AttributeError:
-        # Fallback if the model doesn't support probabilities
-        pred = model.predict(features)[0]
-        # Mock probabilities if exact confidence isn't available
-        pred_prob = [1.0, 0.0] if pred == 0 else [0.0, 1.0]
+        # --- A. Extract Features ---
+        length = len(url)
+        dots = url.count('.')
+        hyphens = url.count('-')
+        digits = sum(c.isdigit() for c in url)
+        at_sign = url.count('@')
+        
+        # Raw features list
+        features_list = [length, dots, hyphens, digits, at_sign]
 
-    # 3. Format Output
-    # Assuming Index 0 = Safe, Index 1 = Phishing
-    label_index = pred_prob.argmax()
-    confidence = float(pred_prob[label_index])
-    
-    if label_index == 1:
-        label = "🚨 Phishing"
-    else:
-        label = "✅ Safe"
-    
-    return {label: confidence}
+        # --- B. Smart Detection & Prediction ---
+        
+        # CASE 1: It is a Scikit-Learn Model (Random Forest, etc.)
+        if hasattr(model, "predict_proba"):
+            # Sklearn expects a Numpy Array
+            input_data = np.array([features_list], dtype=float)
+            
+            pred_prob = model.predict_proba(input_data)[0]
+            # Usually: Index 0 = Safe, Index 1 = Phishing
+            safe_score = float(pred_prob[0])
+            phish_score = float(pred_prob[1])
+
+        # CASE 2: It is a PyTorch Neural Network
+        elif isinstance(model, torch.nn.Module):
+            model.eval() # Set to evaluation mode
+            
+            # PyTorch expects a Tensor
+            input_tensor = torch.tensor([features_list], dtype=torch.float32)
+            
+            with torch.no_grad():
+                logits = model(input_tensor)
+                
+                # Check output shape to decide between Softmax or Sigmoid
+                if logits.shape[1] == 1:
+                    # Binary output (Sigmoid)
+                    phish_score = torch.sigmoid(logits).item()
+                    safe_score = 1.0 - phish_score
+                else:
+                    # Multi-class output (Softmax)
+                    probs = torch.nn.functional.softmax(logits, dim=1)
+                    safe_score = float(probs[0][0])
+                    phish_score = float(probs[0][1])
+        
+        else:
+            return {"Error": 0}, f"Unknown model type: {type(model)}"
+
+        # Return results
+        return {"✅ Safe": safe_score, "🚨 Phishing": phish_score}, "Success"
+
+    except Exception as e:
+        # This catches the specific error and shows it in the UI
+        error_msg = f"Crash Error: {str(e)}"
+        print(error_msg)
+        return {"Error": 0}, error_msg
 
 # -----------------------------
-# 4. Unique Professional UI
+# 3. UI Setup
 # -----------------------------
-theme = gr.themes.Soft(
-    primary_hue="blue",
-    secondary_hue="slate",
-)
-
-with gr.Blocks(theme=theme, title="PhishGuard Local") as iface:
+with gr.Blocks(theme=gr.themes.Soft()) as iface:
+    gr.Markdown("# 🛡️ PhishGuard Debugger")
     
-    # Header
     with gr.Row():
-        gr.Markdown(
-            """
-            # 🛡️ PhishGuard (Local Model)
-            ### Custom Random Forest Detector
-            Running locally using your `phishing_rf_model.pt` file.
-            """
-        )
+        input_box = gr.Textbox(label="URL", placeholder="https://google.com")
+        predict_btn = gr.Button("Scan", variant="primary")
     
-    # Main Interface
     with gr.Row():
-        with gr.Column(scale=1):
-            url_input = gr.Textbox(
-                lines=3, 
-                placeholder="https://example.com", 
-                label="Check URL",
-                info="Paste the link you want to test."
-            )
-            submit_btn = gr.Button("Scan URL 🔍", variant="primary")
-            
-            gr.Examples(
-                examples=[
-                    ["https://google.com"], 
-                    ["http://fake-login-secure.com/update"]
-                ],
-                inputs=url_input
-            )
-
-        with gr.Column(scale=1):
-            output_label = gr.Label(label="Result")
-            gr.Markdown("> **Note:** Running without feature scaler. Results depend on raw feature interpretation.")
+        # We use two outputs: one for the label, one for the error message
+        output_label = gr.Label(label="Prediction")
+        status_box = gr.Textbox(label="Debug Status (Read this if error)", interactive=False)
 
-    # Actions
-    submit_btn.click(
-        fn=predict_phishing,
-        inputs=url_input,
-        outputs=output_label
+    predict_btn.click(
+        fn=predict_phishing, 
+        inputs=input_box, 
+        outputs=[output_label, status_box]
     )
 
-iface.launch(share=True)
+iface.launch()