Upload app.py

app.py

@@ -1,223 +1,54 @@
-import os
-from typing import List, Optional, Dict, Any
-
-import torch
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSequenceClassification,
-    AutoConfig,
-)
-
-# -----------------------------
-# Runtime & cache configuration
-# -----------------------------
-# 兼容容器/无写权限环境：把 HF/torch 缓存指到可写目录
-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")
-
-CACHE_DIR = os.environ.get("HF_CACHE_DIR", "/data/.cache")
-os.makedirs(CACHE_DIR, exist_ok=True)
-
-# -----------------------------
-# Model configuration
-# -----------------------------
-MODEL_ID = os.environ.get("MODEL_ID", "Perth0603/phishing-email-mobilebert")
-MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "256"))  # 和你最初示例保持一致的默认 256
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE_FP16 = bool(int(os.environ.get("FP16", "0"))) and torch.cuda.is_available()
-
-# 限制 CPU 线程数（避免过度并发导致抖动）
-torch.set_num_threads(int(os.environ.get("TORCH_NUM_THREADS", "1")))
-
-# -----------------------------
-# FastAPI app
-# -----------------------------
-app = FastAPI(title="Text Classifier API", version="2.0.0")
-
-# 可按需放开跨域
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=[os.environ.get("CORS_ALLOW_ORIGINS", "*")],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# -----------------------------
-# Pydantic schemas
-# -----------------------------
-class PredictRequest(BaseModel):
-    texts: List[str] = Field(default_factory=list, description="要预测的文本列表")
-    # 可选：如果你只想返回指定阈值的二分类标签，可以传 threshold
-    threshold: Optional[float] = Field(
-        default=None,
-        ge=0.0,
-        le=1.0,
-        description="二分类时用于将正类概率转为标签的阈值（不传则使用 argmax）",
-    )
-
-
-class PredictResult(BaseModel):
-    label: int
-    score: float  # 按你示例，返回的是类别1（正类）的概率
-
-
-class PredictResponse(BaseModel):
-    results: List[PredictResult]
-
-
-# -----------------------------
-# Lazy singletons
-# -----------------------------
-_tokenizer = None
-_model = None
-_id2label: Dict[int, str] = {}
-_label2id: Dict[str, int] = {}
-
-
-def _load_model() -> None:
-    """Lazy load model/tokenizer/config; warmup once."""
-    global _tokenizer, _model, _id2label, _label2id
-    if _tokenizer is not None and _model is not None:
-        return
-
-    config = AutoConfig.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-    _id2label = {int(k): v for k, v in getattr(config, "id2label", {}).items()} or {
-        0: "NEGATIVE",
-        1: "POSITIVE",
-    }
-    _label2id = {v: k for k, v in _id2label.items()}
-
-    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-
-    # dtype 仅在 GPU 下且你明确开启 FP16 时使用
-    dtype = torch.float16 if (DTYPE_FP16 and DEVICE == "cuda") else None
-    _model = AutoModelForSequenceClassification.from_pretrained(
-        MODEL_ID, cache_dir=CACHE_DIR, torch_dtype=dtype
-    ).to(DEVICE)
-
-    # 预热
-    with torch.no_grad():
-        dummy = _tokenizer(["warm up"], return_tensors="pt").to(DEVICE)
-        _ = _model(**dummy).logits
-
-
-# -----------------------------
-# Utility
-# -----------------------------
-def _to_int_label_from_argmax(idx: int) -> int:
-    """
-    把 argmax 的类别索引转成 int 标签。
-    若你的业务希望直接返回索引（如 0/1），这里就保持原样。
-    """
-    return int(idx)
-
-
-def _to_int_label_from_threshold(p_pos: float, threshold: float) -> int:
-    """
-    使用阈值将正类概率转为 0/1 标签。
-    """
-    return int(p_pos >= threshold)
-
-
-# -----------------------------
-# Routes
-# -----------------------------
-@app.get("/")
-def root() -> Dict[str, Any]:
-    return {
-        "status": "ok",
-        "model": MODEL_ID,
-        "device": DEVICE,
-        "max_length": MAX_LENGTH,
-        "version": app.version,
-    }
-
-
-@app.get("/healthz")
-def healthz() -> Dict[str, str]:
-    return {"status": "ok"}
-
-
-@app.get("/readyz")
-def readyz() -> Dict[str, str]:
-    try:
-        _load_model()
-        return {"status": "ready"}
-    except Exception as e:
-        return {"status": "error", "detail": str(e)}
-
-
-@app.post("/predict", response_model=PredictResponse)
-def predict(req: PredictRequest):
-    """
-    批量预测入口：
-    - 输入：texts: List[str]
-    - 输出：每条文本的 {"label": int, "score": float}，其中 score = 概率(类别1)
-    - label 的生成规则：
-        * 若提供 threshold：label = 1 当 概率(1) >= threshold；否则为 0
-        * 不提供 threshold：label = argmax(prob)
-    """
-    if not req.texts:
-        return {"results": []}
-
-    try:
-        _load_model()
-
-        # Tokenize（批量），padding=True + truncation=True，长度可配
-        encoded = _tokenizer(
-            req.texts,
-            padding=True,
-            truncation=True,
-            max_length=MAX_LENGTH,
-            return_tensors="pt",
-        ).to(DEVICE)
-
-        with torch.no_grad():
-            # autocast 仅在 FP16/GPU 开启时使用
-            if DTYPE_FP16 and DEVICE == "cuda":
-                with torch.cuda.amp.autocast(dtype=torch.float16):
-                    logits = _model(**encoded).logits
-            else:
-                logits = _model(**encoded).logits
-
-            probs = torch.softmax(logits, dim=-1).detach().cpu().numpy()
-            preds = probs.argmax(axis=-1)
-
-        results: List[PredictResult] = []
-        for i in range(len(req.texts)):
-            # 二分类：取第1类（正类）的概率作为 score
-            p_pos = float(probs[i, 1]) if probs.shape[1] > 1 else float(probs[i, 0])
-
-            if req.threshold is not None and probs.shape[1] == 2:
-                label = _to_int_label_from_threshold(p_pos, req.threshold)
-            else:
-                label = _to_int_label_from_argmax(int(preds[i]))
-
-            results.append(PredictResult(label=label, score=p_pos))
-
-        return {"results": [r.dict() for r in results]}
-    except Exception as e:
-        # 返回 500 + 错误信息（方便排障）
-        return JSONResponse(status_code=500, content={"error": str(e)})
-
-
-# -----------------------------
-# Local dev entry
-# -----------------------------
-if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run(
-        "app:app",
-        host=os.environ.get("HOST", "0.0.0.0"),
-        port=int(os.environ.get("PORT", "7860")),
-        reload=bool(int(os.environ.get("RELOAD", "0"))),
-    )
+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()
+    try:
+        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 {"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)}
+
+