Update app.py

app.py

@@ -1,50 +1,156 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import os
-
-
-MODEL_ID = os.environ.get("MODEL_ID", "Perth0603/phishing-email-mobilebert")
-
-app = FastAPI(title="Phishing Text Classifier", version="1.0.0")
-
-
-class PredictPayload(BaseModel):
-    inputs: str
-
-
-# Lazy singletons for model/tokenizer
-_tokenizer = None
-_model = None
-
-
-def _load_model():
-    global _tokenizer, _model
-    if _tokenizer is None or _model is None:
-        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
-        # Warm-up
-        with torch.no_grad():
-            _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
-
-
-@app.get("/")
-def root():
-    return {"status": "ok", "model": MODEL_ID}
-
-
-@app.post("/predict")
-def predict(payload: PredictPayload):
-    _load_model()
-    with torch.no_grad():
-        logits = _model(**_tokenizer([payload.inputs], return_tensors="pt")).logits
-        probs = torch.softmax(logits, dim=-1)[0]
-        score, idx = torch.max(probs, dim=0)
-
-    # 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)}
-
-
+import os
+
+os.environ.setdefault("HOME", "/data")
+os.environ.setdefault("XDG_CACHE_HOME", "/data/.cache")
+os.environ.setdefault("HF_HOME", "/data/.cache")
+os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.cache")
+os.environ.setdefault("TORCH_HOME", "/data/.cache")
+
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+MODEL_ID = os.environ.get("MODEL_ID", "Perth0603/phishing-email-mobilebert")
+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.2.0")
+
+class PredictPayload(BaseModel):
+    inputs: str
+
+_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
+
+
+@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.",
+    }
+
+
+@app.post("/predict")
+def predict(payload: PredictPayload, debug: bool = False):
+    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
+
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})