Update app.py

app.py

@@ -2,17 +2,21 @@ import gradio as gr
 import numpy as np
 import tensorflow as tf
 import joblib
-import requests
+import torch
 import os
 
+from concurrent.futures import ThreadPoolExecutor
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 from tensorflow.keras.layers import Layer
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
 
 print("Loading models...")
 
-############################################
+
+########################################
 # Custom Attention Layer
-############################################
+########################################
 
 class AttentionLayer(Layer):
 
@@ -37,15 +41,20 @@ class AttentionLayer(Layer):
         return context
 
 
-############################################
-# Load Models
-############################################
+########################################
+# Load Classical ML Models
+########################################
 
 vectorizer = joblib.load("models/tfidf_vectorizer.pkl")
 
 lr_model = joblib.load("models/logistic_model.pkl")
 rf_model = joblib.load("models/random_forest_model.pkl")
 
+
+########################################
+# Load Tokenizer + Neural Models
+########################################
+
 tokenizer = joblib.load("models/tokenizer.pkl")
 
 lstm_model = tf.keras.models.load_model(
@@ -57,28 +66,32 @@ bilstm_model = tf.keras.models.load_model(
     custom_objects={"AttentionLayer": AttentionLayer}
 )
 
-print("All models loaded successfully")
 
+########################################
+# Load DistilBERT
+########################################
 
-############################################
-# HuggingFace DistilBERT Setup
-############################################
+print("Loading DistilBERT...")
 
-BERT_API = "https://router.huggingface.co/hf-inference/pipeline/text-classification/Savizzz-35/phishguard-distilbert"
+bert_tokenizer = AutoTokenizer.from_pretrained(
+    "Savizzz-35/phishguard-distilbert"
+)
+
+bert_model = AutoModelForSequenceClassification.from_pretrained(
+    "Savizzz-35/phishguard-distilbert"
+)
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+bert_model.eval()
 
-print("HF TOKEN LOADED:", HF_TOKEN is not None)
+print("DistilBERT loaded")
 
-headers = {
-    "Authorization": f"Bearer {HF_TOKEN}",
-    "Content-Type": "application/json"
-}
+
+print("All models loaded successfully")
 
 
-############################################
-# Feature Fix for sklearn models
-############################################
+########################################
+# Fix feature mismatch (sklearn)
+########################################
 
 def adjust_features(vec, model):
 
@@ -97,9 +110,9 @@ def adjust_features(vec, model):
     return vec
 
 
-############################################
+########################################
 # Individual Model Predictions
-############################################
+########################################
 
 def lr_predict(text):
 
@@ -141,64 +154,29 @@ def bilstm_predict(text):
     return float(prob)
 
 
-############################################
-# DEBUG DISTILBERT
-############################################
-
 def bert_predict(text):
 
-    payload = {
-        "inputs": text
-    }
-
-    print("\n============================")
-    print("BERT DEBUG START")
-    print("Payload:", payload)
-
-    try:
-
-        r = requests.post(
-            BERT_API,
-            headers=headers,
-            json=payload,
-            timeout=30
-        )
-
-        print("HTTP STATUS:", r.status_code)
-        print("RAW RESPONSE:", r.text)
-
-        result = r.json()
-
-        print("JSON RESPONSE:", result)
-
-        if isinstance(result, list):
-
-            if len(result) > 0 and isinstance(result[0], list):
-                result = result[0]
-
-            for item in result:
-
-                if item["label"] == "phishing_email":
-
-                    score = float(item["score"])
-                    print("BERT SCORE:", score)
-                    print("============================\n")
-
-                    return score
+    inputs = bert_tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    )
 
-    except Exception as e:
+    with torch.no_grad():
 
-        print("BERT EXCEPTION:", str(e))
+        outputs = bert_model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=1)
 
-    print("BERT FAILED → returning fallback 0.5")
-    print("============================\n")
+        phishing_prob = probs[0][1].item()
 
-    return 0.5
+    return float(phishing_prob)
 
 
-############################################
-# Ensemble
-############################################
+########################################
+# Ensemble Weights
+########################################
 
 WEIGHTS = {
     "lr": 0.25,
@@ -209,20 +187,31 @@ WEIGHTS = {
 }
 
 
+########################################
+# Parallel Prediction
+########################################
+
 def predict_email(text):
 
-    lr = lr_predict(text)
-    rf = rf_predict(text)
-    lstm = lstm_predict(text)
-    bilstm = bilstm_predict(text)
-    bert = bert_predict(text)
+    with ThreadPoolExecutor() as executor:
+
+        futures = {
+
+            "lr": executor.submit(lr_predict, text),
+            "rf": executor.submit(rf_predict, text),
+            "lstm": executor.submit(lstm_predict, text),
+            "bilstm": executor.submit(bilstm_predict, text),
+            "bert": executor.submit(bert_predict, text)
+        }
+
+        results = {name: f.result() for name, f in futures.items()}
 
     final_score = (
-        lr * WEIGHTS["lr"]
-        + rf * WEIGHTS["rf"]
-        + lstm * WEIGHTS["lstm"]
-        + bilstm * WEIGHTS["bilstm"]
-        + bert * WEIGHTS["bert"]
+        results["lr"] * WEIGHTS["lr"]
+        + results["rf"] * WEIGHTS["rf"]
+        + results["lstm"] * WEIGHTS["lstm"]
+        + results["bilstm"] * WEIGHTS["bilstm"]
+        + results["bert"] * WEIGHTS["bert"]
     )
 
     if final_score < 0.5:
@@ -235,19 +224,13 @@ def predict_email(text):
     return {
         "phishing_probability": float(final_score),
         "risk_level": risk,
-        "model_outputs": {
-            "lr": lr,
-            "rf": rf,
-            "lstm": lstm,
-            "bilstm": bilstm,
-            "bert": bert
-        }
+        "model_outputs": results
     }
 
 
-############################################
-# API Wrapper
-############################################
+########################################
+# API wrapper
+########################################
 
 def run_detection(text):
 
@@ -260,9 +243,9 @@ def run_detection(text):
         return {"error": str(e)}
 
 
-############################################
-# Gradio UI
-############################################
+########################################
+# Gradio Interface
+########################################
 
 iface = gr.Interface(
     fn=run_detection,