Upload 4 files

app.py

@@ -12,12 +12,20 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from huggingface_hub import hf_hub_download
 import joblib
 import torch
+import re
+import numpy as np
+import pandas as pd
+try:
+    import xgboost as xgb  # type: ignore
+except Exception:
+    xgb = None  # optional; required if bundle uses xgboost
 
 
 MODEL_ID = os.environ.get("MODEL_ID", "Perth0603/phishing-email-mobilebert")
 URL_REPO = os.environ.get("URL_REPO", "Perth0603/Random-Forest-Model-for-PhishingDetection")
 URL_REPO_TYPE = os.environ.get("URL_REPO_TYPE", "model")  # model|space|dataset
-URL_FILENAME = os.environ.get("URL_FILENAME", "url_rf_model.joblib")
+# NOTE: set to your artifact filename, e.g. rf_url_phishing_xgboost_bst.joblib
+URL_FILENAME = os.environ.get("URL_FILENAME", "rf_url_phishing_xgboost_bst.joblib")
 
 # Ensure writable cache directory for HF/torch inside Spaces Docker
 CACHE_DIR = os.environ.get("HF_CACHE_DIR", "/data/.cache")
@@ -33,16 +41,16 @@ class PredictPayload(BaseModel):
 # Lazy singletons for model/tokenizer
 _tokenizer = None
 _model = None
-_url_model = None
+_url_bundle = None  # holds dict: {model, feature_cols, url_col, label_col, model_type}
 
 
 def _load_url_model():
-    global _url_model
-    if _url_model is None:
+    global _url_bundle
+    if _url_bundle is None:
         # Prefer local artifact if present (e.g., committed into the Space repo)
         local_path = os.path.join(os.getcwd(), URL_FILENAME)
         if os.path.exists(local_path):
-            _url_model = joblib.load(local_path)
+            _url_bundle = joblib.load(local_path)
             return
         # Download model artifact from HF Hub
         model_path = hf_hub_download(
@@ -51,7 +59,32 @@ def _load_url_model():
             repo_type=URL_REPO_TYPE,
             cache_dir=CACHE_DIR,
         )
-        _url_model = joblib.load(model_path)
+        _url_bundle = joblib.load(model_path)
+
+
+# URL feature engineering (must match training)
+_SUSPICIOUS_TOKENS = ["login", "verify", "secure", "update", "bank", "pay", "account", "webscr"]
+_ipv4_pattern = re.compile(r'(?:\d{1,3}\.){3}\d{1,3}')
+
+def _engineer_features(df: pd.DataFrame, url_col: str, feature_cols: list[str] | None = None) -> pd.DataFrame:
+    s = df[url_col].astype(str)
+    out = pd.DataFrame(index=df.index)
+    out['url_len'] = s.str.len().fillna(0)
+    out['count_dot'] = s.str.count(r'\.')
+    out['count_hyphen'] = s.str.count('-')
+    out['count_digit'] = s.str.count(r'\d')
+    out['count_at'] = s.str.count('@')
+    out['count_qmark'] = s.str.count('\?')
+    out['count_eq'] = s.str.count('=')
+    out['count_slash'] = s.str.count('/')
+    out['digit_ratio'] = (out['count_digit'] / out['url_len'].replace(0, np.nan)).fillna(0)
+    out['has_ip'] = s.str.contains(_ipv4_pattern).astype(int)
+    for tok in _SUSPICIOUS_TOKENS:
+        out[f'has_{tok}'] = s.str.contains(tok, case=False, regex=False).astype(int)
+    out['starts_https'] = s.str.startswith('https').astype(int)
+    out['ends_with_exe'] = s.str.endswith('.exe').astype(int)
+    out['ends_with_zip'] = s.str.endswith('.zip').astype(int)
+    return out if feature_cols is None else out[feature_cols]
 
 
 def _load_model():
@@ -95,26 +128,39 @@ class PredictUrlPayload(BaseModel):
 def predict_url(payload: PredictUrlPayload):
     try:
         _load_url_model()
-        # Expect the sklearn pipeline to accept raw URL string and output either
-        # probabilities via predict_proba or binary via predict
-        model = _url_model
+        bundle = _url_bundle
+        if not isinstance(bundle, dict) or 'model' not in bundle:
+            raise RuntimeError("Loaded URL artifact is not a bundle dict with 'model'.")
+        model = bundle['model']
+        feature_cols = bundle.get('feature_cols') or []
+        url_col = bundle.get('url_col') or 'url'
+        model_type = bundle.get('model_type') or ''
+
+        row = pd.DataFrame({url_col: [payload.url]})
+        feats = _engineer_features(row, url_col, feature_cols)
+
         score = None
         label = None
-        if hasattr(model, "predict_proba"):
-            proba = model.predict_proba([payload.url])[0]
-            # Assume index 1 corresponds to PHISH
+
+        if isinstance(model_type, str) and model_type == 'xgboost_bst':
+            if xgb is None:
+                raise RuntimeError("xgboost is not installed but required for this model bundle.")
+            dmat = xgb.DMatrix(feats)
+            proba = float(model.predict(dmat)[0])
+            score = proba
+            label = "PHISH" if score >= 0.5 else "LEGIT"
+        elif hasattr(model, "predict_proba"):
+            proba = model.predict_proba(feats)[0]
             if len(proba) == 2:
                 score = float(proba[1])
                 label = "PHISH" if score >= 0.5 else "LEGIT"
             else:
-                # Multiclass fallback: take max
-                max_idx = int(proba.argmax())
+                max_idx = int(np.argmax(proba))
                 score = float(proba[max_idx])
                 label = "PHISH" if max_idx == 1 else "LEGIT"
         else:
-            pred = model.predict([payload.url])[0]
-            # Common encodings: 1=phish, 0=legit or strings
-            if isinstance(pred, (int, float)):
+            pred = model.predict(feats)[0]
+            if isinstance(pred, (int, float, np.integer, np.floating)):
                 label = "PHISH" if int(pred) == 1 else "LEGIT"
                 score = 1.0 if label == "PHISH" else 0.0
             else:

requirements.txt

@@ -10,4 +10,6 @@ safetensors>=0.4.3
 huggingface_hub>=0.23.0
 scikit-learn>=1.3.0
 joblib>=1.3.0
+pandas>=2.0.0
+xgboost>=2.0.0