Upload app.py

app.py

@@ -65,6 +65,27 @@ _url_lock = threading.Lock()
 # Calibrated flag: is XGB class 1 == PHISH?
 _url_phish_is_positive: Optional[bool] = None
 
+# -------------------------
+# Autocalibration URL prototypes (editable)
+# -------------------------
+# You can edit these lists to define which URLs are considered obviously phishy/legit
+# for polarity auto-calibration of classical URL models (e.g., XGBoost, scikit-learn).
+_AUTOCALIB_PHISHY_URLS: List[str] = [
+    "http://198.51.100.23/login/update?acc=123",
+    "http://secure-login-account-update.example.com/session?id=123",
+    "http://bank.verify-update-security.com/confirm",
+    "http://paypal.com.account-verify.cn/login",
+    "http://abc.xyz/downloads/invoice.exe",
+]
+
+_AUTOCALIB_LEGIT_URLS: List[str] = [
+    "https://www.wikipedia.org/",
+    "https://www.microsoft.com/",
+    "https://www.openai.com/",
+    "https://www.python.org/",
+    "https://www.gov.uk/",
+]
+
 # -------------------------
 # URL features (must match training)
 # -------------------------
@@ -95,7 +116,7 @@ def _engineer_features(df: pd.DataFrame, url_col: str, feature_cols: Optional[Li
 # Loaders
 # -------------------------
 def _load_model():
-    global _tokenizer, _model, _id2label, _label2id
+    global _tokenizer, _model, _id2label, _label2id, _text_phish_id
     if _tokenizer is None or _model is None:
         with _model_lock:
             if _tokenizer is None or _model is None:
@@ -105,9 +126,26 @@ def _load_model():
                 if cfg is not None and getattr(cfg, "id2label", None):
                     _id2label = {int(k): v for k, v in cfg.id2label.items()}
                     _label2id = {v: int(k) for k, v in _id2label.items()}
+                # Try to detect which index corresponds to PHISH/SPAM
+                _text_phish_id = _detect_text_phish_id(_id2label)
                 with torch.no_grad():
                     _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
 
+# Detect which label id corresponds to phishing for text models based on label strings
+_text_phish_id: Optional[int] = None
+
+def _detect_text_phish_id(id2label: Dict[int, str]) -> Optional[int]:
+    candidates_phish = ("PHISH", "SPAM", "MALICIOUS", "POSITIVE")
+    # Prefer explicit PHISH/SPAM over generic POSITIVE
+    priority_order = ("PHISH", "SPAM", "MALICIOUS", "POSITIVE")
+    norm = {k: str(v).strip().upper() for k, v in id2label.items()}
+    # exact/substring match in priority order
+    for token in priority_order:
+        for k, v in norm.items():
+            if token in v:
+                return int(k)
+    return None
+
 def _load_url_model():
     global _url_bundle
     if _url_bundle is None:
@@ -144,20 +182,8 @@ def _auto_calibrate_phish_positive(bundle: Dict[str, Any], feature_cols: List[st
     if "phish_is_positive" in bundle:
         return bool(bundle["phish_is_positive"])
 
-    phishy = [
-        "http://198.51.100.23/login/update?acc=123",
-        "http://secure-login-account-update.example.com/session?id=123",
-        "http://bank.verify-update-security.com/confirm",
-        "http://paypal.com.account-verify.cn/login",
-        "http://abc.xyz/downloads/invoice.exe"
-    ]
-    legit = [
-        "https://www.wikipedia.org/",
-        "https://www.microsoft.com/",
-        "https://www.openai.com/",
-        "https://www.python.org/",
-        "https://www.gov.uk/"
-    ]
+    phishy = _AUTOCALIB_PHISHY_URLS
+    legit = _AUTOCALIB_LEGIT_URLS
 
     def _batch_mean(urls: List[str]) -> float:
         df = pd.DataFrame({url_col: urls})
@@ -218,8 +244,17 @@ def predict(payload: PredictPayload):
             logits = _model(**inputs).logits
             probs = torch.softmax(logits, dim=-1)[0]
             score, idx = torch.max(probs, dim=0)
-            label = _id2label.get(int(idx), str(int(idx)))
-        return {"label": label, "score": float(score), "raw_index": int(idx)}
+
+            # Normalize label to PHISH/LEGIT if we could detect PHISH id
+            if _text_phish_id is not None and 0 <= _text_phish_id < probs.shape[0]:
+                phish_prob = float(probs[_text_phish_id])
+                norm_label = "PHISH" if phish_prob >= 0.5 else "LEGIT"
+                norm_score = phish_prob if norm_label == "PHISH" else (1.0 - phish_prob)
+                return {"label": norm_label, "score": float(norm_score), "raw_index": int(idx)}
+            else:
+                # Fallback to model's provided labels
+                label = _id2label.get(int(idx), str(int(idx)))
+                return {"label": label, "score": float(score), "raw_index": int(idx)}
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})