Upload app.py

app.py

@@ -1,5 +1,4 @@
 import os
-
 os.environ.setdefault("HOME", "/data")
 os.environ.setdefault("XDG_CACHE_HOME", "/data/.cache")
 os.environ.setdefault("HF_HOME", "/data/.cache")
@@ -12,145 +11,55 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
+
 MODEL_ID = os.environ.get("MODEL_ID", "Perth0603/phishing-email-mobilebert")
+
+# Ensure writable cache directory for HF/torch inside Spaces Docker
 CACHE_DIR = os.environ.get("HF_CACHE_DIR", "/data/.cache")
 os.makedirs(CACHE_DIR, exist_ok=True)
 
-# Optional overrides:
-# - set PHISH_INDEX=0 or 1 to force which logit is "phish"
-# - or POSITIVE_CLASS=PHISH/LEGIT to force polarity if the head is single-logit
-PHISH_INDEX_ENV = os.environ.get("PHISH_INDEX", "").strip()
-POS_CLASS_ENV = os.environ.get("POSITIVE_CLASS", "").strip().upper()  # for single-logit heads
+app = FastAPI(title="Phishing Text Classifier", version="1.0.0")
 
-app = FastAPI(title="Phishing Text Classifier", version="1.2.0")
 
 class PredictPayload(BaseModel):
     inputs: str
 
+
+# Lazy singletons for model/tokenizer
 _tokenizer = None
 _model = None
-_id2label = None
-_phish_idx = None  # which class id means "phish" for multi-logit heads
-
-
-def _norm(lbl: str) -> str:
-    return (lbl or "").lower().replace("label_", "").replace("_", "").strip()
-
-
-def _auto_pick_phish_idx(id2label: dict) -> int | None:
-    # Try to find an explicit phishing-like label
-    normalized = {int(k): _norm(v) for k, v in id2label.items()}
-    for i, s in normalized.items():
-        if any(k in s for k in ("phish", "phishing", "spam", "malicious")):
-            return int(i)
-    # Try to infer from legit/ham/safe presence
-    legit_like = [i for i, s in normalized.items() if any(k in s for k in ("legit", "ham", "safe", "nonphish"))]
-    if len(legit_like) == 1 and len(id2label) == 2:
-        return 1 - int(legit_like[0])
-    return None  # unknown
-
-
-def _calibrate_phish_idx(model, tokenizer, num_labels: int) -> int:
-    # Probe with typical phishing vs legitimate texts to decide which index behaves like "phish"
-    phish_samples = [
-        "Urgent: Your account has been suspended. Verify your password at http://login-security-update.com",
-        "We detected unusual activity. Confirm your identity now or your account will be closed.",
-    ]
-    legit_samples = [
-        "Let's schedule a meeting for tomorrow to discuss the project status.",
-        "Thank you for your purchase. Here is your receipt for your recent order.",
-    ]
-    with torch.no_grad():
-        p_logits = model(**tokenizer(phish_samples, return_tensors="pt", padding=True, truncation=True, max_length=512)).logits
-        l_logits = model(**tokenizer(legit_samples, return_tensors="pt", padding=True, truncation=True, max_length=512)).logits
-        if p_logits.shape[-1] == 1:
-            # Single-logit: positive class is sigmoid(logit). If POSITIVE_CLASS env says PHISH, use it; else assume PHISH
-            return 0  # index unused for single-logit; return anything
-        probs_p = torch.softmax(p_logits, dim=-1).mean(dim=0)  # [C]
-        probs_l = torch.softmax(l_logits, dim=-1).mean(dim=0)  # [C]
-        # Pick the index that is high for phishing samples and low for legit samples
-        diffs = probs_p - probs_l
-        return int(torch.argmax(diffs).item())
 
 
 def _load_model():
-    global _tokenizer, _model, _id2label, _phish_idx
-    if _tokenizer is not None and _model is not None:
-        return
-
-    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-    _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-    _model.eval()
-
-    cfg = _model.config
-    id2label = getattr(cfg, "id2label", None) or {i: f"LABEL_{i}" for i in range(int(getattr(cfg, "num_labels", 2)))}
-    _id2label = {int(k): str(v) for k, v in id2label.items()}
-    num_labels = len(_id2label)
-
-    # 1) Respect explicit env override
-    if PHISH_INDEX_ENV in ("0", "1"):
-        _phish_idx = int(PHISH_INDEX_ENV)
-    else:
-        # 2) Try to infer from label names
-        auto_idx = _auto_pick_phish_idx(_id2label) if num_labels >= 2 else None
-        if auto_idx is not None:
-            _phish_idx = auto_idx
-        else:
-            # 3) Calibrate by probing example texts (multi-logit only)
-            _phish_idx = _calibrate_phish_idx(_model, _tokenizer, num_labels) if num_labels >= 2 else 0
-
-    # Warm-up
-    with torch.no_grad():
-        _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
+    global _tokenizer, _model
+    if _tokenizer is None or _model is None:
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
+        _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
+        # Warm-up
+        with torch.no_grad():
+            _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
 
 
 @app.get("/")
 def root():
-    return {
-        "status": "ok",
-        "model": MODEL_ID,
-        "id2label": _id2label,
-        "phish_index": _phish_idx,
-        "note": "Set PHISH_INDEX=0|1 to override mapping if needed.",
-    }
+    return {"status": "ok", "model": MODEL_ID}
 
 
 @app.post("/predict")
-def predict(payload: PredictPayload, debug: bool = False):
+def predict(payload: PredictPayload):
     try:
         _load_model()
         with torch.no_grad():
             inputs = _tokenizer([payload.inputs], return_tensors="pt", truncation=True, max_length=512)
-            logits = _model(**inputs).logits  # [1, C] or [1, 1]
-
-            if logits.shape[-1] == 1:
-                # Single-logit: sigmoid gives probability of the model's POSITIVE class.
-                p_pos = torch.sigmoid(logits[0, 0]).item()
-                # If POSITIVE_CLASS says PHISH, treat positive as PHISH; else treat positive as LEGIT.
-                positive_is_phish = (POS_CLASS_ENV == "PHISH") or (POS_CLASS_ENV == "")
-
-                phish_proba = p_pos if positive_is_phish else (1.0 - p_pos)
-                probs = torch.tensor([1.0 - phish_proba, phish_proba])  # for debugging
-            else:
-                probs = torch.softmax(logits, dim=-1)[0]
-                phish_proba = float(probs[int(_phish_idx)])
-
-            label = "PHISH" if phish_proba >= 0.5 else "LEGIT"
-            frontend_label = "Phishing" if label == "PHISH" else "Legitimate"
-
-            out = {
-                "label": label,
-                "frontend_label": frontend_label,
-                "phishing_probability": float(phish_proba),
-            }
-            if debug:
-                out.update({
-                    "id2label": _id2label,
-                    "phish_index": int(_phish_idx),
-                    "logits": [float(x) for x in logits[0]],
-                    "probs": [float(x) for x in (probs if isinstance(probs, torch.Tensor) else torch.tensor(probs))],
-                })
-            return out
-
+            logits = _model(**inputs).logits
+            probs = torch.softmax(logits, dim=-1)[0]
+            score, idx = torch.max(probs, dim=0)
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": str(e)})
+
+    # Map common ids to labels (kept generic; your config also has these)
+    id2label = {0: "LEGIT", 1: "PHISH"}
+    label = id2label.get(int(idx), str(int(idx)))
+    return {"label": label, "score": float(score)}
+
+