smaaanwerb/postsmoderator

1import os2import torch3from fastapi import FastAPI, HTTPException4from fastapi.middleware.cors import CORSMiddleware5from pydantic import BaseModel6from transformers import AutoTokenizer, AutoModelForSequenceClassification7from huggingface_hub import login8import uvicorn9 10# ─────────────────────────────────────────11#  Config12# ─────────────────────────────────────────13MODEL_NAME = os.getenv("MODEL_NAME", "smaaanwerb/posts")  14HF_TOKEN   = os.getenv("HF_TOKEN")   # ← Secret في إعدادات الـ Space15DEVICE     = torch.device("cuda" if torch.cuda.is_available() else "cpu")16 17# تسجيل الدخول لو الموديل private18if not HF_TOKEN:19    raise RuntimeError("HF_TOKEN is not set — required for private model access")20 21print(f"HF_TOKEN detected, length={len(HF_TOKEN)}, prefix={HF_TOKEN[:6]}...")22 23login(token=HF_TOKEN)24print("Logged in to Hugging Face ✓")25 26LABEL_MAP = {0: "Clean", 1: "Offensive", 2: "Toxic"}27 28# action بيحدد إيه اللي يعمله الـ backend29ACTION_MAP = {30    "Clean":     "publish",   # ينشر عادي31    "Offensive": "publish",   # ينشر بس يسجله32    "Toxic":     "delete",    # يمسح من الـ app بس يحفظه في DB33}34 35# ─────────────────────────────────────────36#  Load Model (عند بداية التطبيق)37# ─────────────────────────────────────────38print(f"Loading model from: {MODEL_NAME}")39 40# نمرر التوكن صراحة بدل ما نعتمد على login() لوحدها41tokenizer = AutoTokenizer.from_pretrained(42    MODEL_NAME,43    use_fast=True,44    token=HF_TOKEN,45)46model = AutoModelForSequenceClassification.from_pretrained(47    MODEL_NAME,48    token=HF_TOKEN,49)50model.to(DEVICE)51model.eval()52print(f"Model loaded on {DEVICE}")53 54# ─────────────────────────────────────────55#  FastAPI App56# ─────────────────────────────────────────57app = FastAPI(58    title="Gym Comment Classifier",59    description="Classifies Arabic/English gym comments into Clean / Offensive / Toxic",60    version="1.0.0",61)62 63app.add_middleware(64    CORSMiddleware,65    allow_origins=["*"],   # ضيّق ده في Production66    allow_methods=["*"],67    allow_headers=["*"],68)69 70# ─────────────────────────────────────────71#  Schemas72# ─────────────────────────────────────────73class ClassifyRequest(BaseModel):74    text: str75    post_id: str | None = None    # اختياري — الـ backend بيبعته لو عايز76 77class ClassifyResponse(BaseModel):78    text:        str79    post_id:     str | None80    label:       str              # Clean | Offensive | Toxic81    action:      str              # publish | delete82    confidence:  dict             # {"Clean": 0.9, "Offensive": 0.07, "Toxic": 0.03}83 84# ─────────────────────────────────────────85#  Core Prediction Function86# ─────────────────────────────────────────87def predict(text: str) -> tuple[str, dict]:88    inputs = tokenizer(89        text,90        return_tensors="pt",91        truncation=True,92        max_length=256,93        padding=True,94    ).to(DEVICE)95 96    with torch.no_grad():97        logits = model(**inputs).logits98        probs  = torch.softmax(logits, dim=1)[0]99        pred   = torch.argmax(probs).item()100 101    label      = LABEL_MAP[pred]102    confidence = {LABEL_MAP[i]: round(probs[i].item(), 4) for i in range(3)}103    return label, confidence104 105# ─────────────────────────────────────────106#  Endpoints107# ─────────────────────────────────────────108@app.get("/")109def root():110    return {"status": "ok", "message": "Gym Comment Classifier is running"}111 112@app.get("/health")113def health():114    return {"status": "healthy", "device": str(DEVICE)}115 116@app.post("/classify", response_model=ClassifyResponse)117def classify(req: ClassifyRequest):118    if not req.text or not req.text.strip():119        raise HTTPException(status_code=422, detail="text field cannot be empty")120 121    try:122        label, confidence = predict(req.text.strip())123    except Exception as e:124        raise HTTPException(status_code=500, detail=f"Model error: {str(e)}")125 126    return ClassifyResponse(127        text=req.text,128        post_id=req.post_id,129        label=label,130        action=ACTION_MAP[label],131        confidence=confidence,132    )133 134# ─────────────────────────────────────────135#  Batch Endpoint (اختياري — لو عندك كتير)136# ─────────────────────────────────────────137class BatchRequest(BaseModel):138    items: list[ClassifyRequest]139 140class BatchResponse(BaseModel):141    results: list[ClassifyResponse]142 143@app.post("/classify/batch", response_model=BatchResponse)144def classify_batch(req: BatchRequest):145    if not req.items:146        raise HTTPException(status_code=422, detail="items list cannot be empty")147    if len(req.items) > 50:148        raise HTTPException(status_code=422, detail="Max 50 items per batch")149 150    results = []151    for item in req.items:152        if not item.text or not item.text.strip():153            continue154        try:155            label, confidence = predict(item.text.strip())156            results.append(ClassifyResponse(157                text=item.text,158                post_id=item.post_id,159                label=label,160                action=ACTION_MAP[label],161                confidence=confidence,162            ))163        except Exception:164            continue165 166    return BatchResponse(results=results)167 168 169if __name__ == "__main__":170    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)