Upload 4 files

README.md

@@ -6,7 +6,10 @@ sdk: docker
 
 # Hugging Face Space - Phishing Text Classifier (Docker + FastAPI)
 
-This Space exposes a minimal `/predict` endpoint for your MobileBERT phishing model so the Flutter app can call it reliably.
+This Space exposes two endpoints so the Flutter app can call them reliably:
+
+- `/predict` for text/email/SMS classification via Transformers
+- `/predict-url` for URL classification via your scikit-learn Random Forest model
 
 ## Files
 - Dockerfile - builds a small FastAPI server image
@@ -18,9 +21,12 @@ This Space exposes a minimal `/predict` endpoint for your MobileBERT phishing mo
 2. Upload the contents of this `hf_space/` folder to the Space root (including Dockerfile).
 3. In Space Settings → Variables, add:
    - MODEL_ID = Perth0603/phishing-email-mobilebert
+   - URL_REPO = Perth0603/Random-Forest-Model-for-PhishingDetection
+   - URL_FILENAME = url_rf_model.joblib  (set to your artifact filename)
 4. Wait for the Space to build and become green. Test:
    - GET `/` should return `{ status: ok, model: ... }`
    - POST `/predict` with `{ "inputs": "Win an iPhone! Click here" }`
+   - POST `/predict-url` with `{ "url": "https://example.com/login" }`
 
 ## Flutter app config
 Set the Space URL in your env file so the app targets the Space instead of the Hosted Inference API:

app.py

@@ -9,10 +9,23 @@ from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 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
+# 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")
@@ -28,6 +41,50 @@ class PredictPayload(BaseModel):
 # Lazy singletons for model/tokenizer
 _tokenizer = None
 _model = None
+_url_bundle = None  # holds dict: {model, feature_cols, url_col, label_col, model_type}
+
+
+def _load_url_model():
+    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_bundle = joblib.load(local_path)
+            return
+        # Download model artifact from HF Hub
+        model_path = hf_hub_download(
+            repo_id=URL_REPO,
+            filename=URL_FILENAME,
+            repo_type=URL_REPO_TYPE,
+            cache_dir=CACHE_DIR,
+        )
+        _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():
@@ -63,3 +120,58 @@ def predict(payload: PredictPayload):
     return {"label": label, "score": float(score)}
 
 
+class PredictUrlPayload(BaseModel):
+    url: str
+
+
+@app.post("/predict-url")
+def predict_url(payload: PredictUrlPayload):
+    try:
+        _load_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 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:
+                max_idx = int(np.argmax(proba))
+                score = float(proba[max_idx])
+                label = "PHISH" if max_idx == 1 else "LEGIT"
+        else:
+            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:
+                up = str(pred).strip().upper()
+                if up in ("PHISH", "PHISHING", "MALICIOUS"):
+                    label, score = "PHISH", 1.0
+                else:
+                    label, score = "LEGIT", 0.0
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
+
+    return {"label": label, "score": float(score)}
+
+

requirements.txt

@@ -6,3 +6,10 @@ torch==2.3.1+cpu
 accelerate>=0.33.0
 safetensors>=0.4.3
 
+# URL model dependencies
+huggingface_hub>=0.23.0
+scikit-learn>=1.3.0
+joblib>=1.3.0
+pandas>=2.0.0
+xgboost>=2.0.0
+