Upload 4 files

Dockerfile

@@ -1,28 +1,28 @@
-FROM python:3.10-slim
-
-ENV PYTHONDONTWRITEBYTECODE=1 \
-    PYTHONUNBUFFERED=1 \
-    PIP_NO_CACHE_DIR=1
-
-WORKDIR /app
-
-# Writable cache directory for HF/torch
-RUN mkdir -p /data/.cache && chmod -R 777 /data
-ENV HF_HOME=/data/.cache \
-    TRANSFORMERS_CACHE=/data/.cache \
-    TORCH_HOME=/data/.cache
-
-# System deps (optional but helps with torch wheels)
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential git && \
-    rm -rf /var/lib/apt/lists/*
-
-COPY requirements.txt /app/requirements.txt
-RUN pip install -r /app/requirements.txt
-
-COPY app.py /app/app.py
-
-EXPOSE 7860
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
-
-
+FROM python:3.10-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+WORKDIR /app
+
+# Writable cache directory for HF/torch
+RUN mkdir -p /data/.cache && chmod -R 777 /data
+ENV HF_HOME=/data/.cache \
+    TRANSFORMERS_CACHE=/data/.cache \
+    TORCH_HOME=/data/.cache
+
+# System deps (optional but helps with torch wheels)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential git && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt /app/requirements.txt
+RUN pip install -r /app/requirements.txt
+
+COPY app.py /app/app.py
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+

app.py

@@ -1,71 +1,66 @@
-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")
-# Optional temperature to sharpen probabilities (lower <1 sharper, >1 smoother)
-TEMP = float(os.environ.get("TEMP", "0.7"))
-
-# 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)
-
-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, cache_dir=CACHE_DIR)
-        _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-        _model.eval()
-        # 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):
-    try:
-        _load_model()
-        with torch.no_grad():
-            inputs = _tokenizer([payload.inputs], return_tensors="pt", truncation=True, max_length=512)
-            logits = _model(**inputs).logits
-            # Raw probs (for reference)
-            raw_probs = torch.softmax(logits, dim=-1)[0]
-            # Temperature-scaled probs to make confidence less around 0.5
-            scaled_probs = torch.softmax(logits / TEMP, dim=-1)[0]
-            score, idx = torch.max(scaled_probs, dim=0)
-    except Exception as 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), "raw_score": float(raw_probs[int(idx)])}
-
-
+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")
+
+# 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)
+
+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, cache_dir=CACHE_DIR)
+        _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
+        _model.eval()
+        # 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):
+    try:
+        _load_model()
+        with torch.no_grad():
+            inputs = _tokenizer([payload.inputs], return_tensors="pt", truncation=True, max_length=512)
+            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)})
+
+    # 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)}
+
+

requirements.txt

@@ -1,8 +1,8 @@
---extra-index-url https://download.pytorch.org/whl/cpu
-fastapi==0.115.0
-uvicorn==0.30.6
-transformers==4.46.3
-torch==2.3.1+cpu
-accelerate>=0.33.0
-safetensors>=0.4.3
-
+--extra-index-url https://download.pytorch.org/whl/cpu
+fastapi==0.115.0
+uvicorn==0.30.6
+transformers==4.46.3
+torch==2.3.1+cpu
+accelerate>=0.33.0
+safetensors>=0.4.3
+