Update app.py

app.py

@@ -16,16 +16,16 @@ from fastapi.middleware.cors import CORSMiddleware
 from jose import jwt, JWTError
 from pydantic import BaseModel, Field
 
-# ── torch shim ─────────────────────────────────────────────────────────
+# ─────────────────────────  torch shim  ───────────────────────────────
 if hasattr(torch, "compile"):
     torch.compile = (lambda m=None,*_,**__: m if callable(m) else (lambda f: f))  # type: ignore
     os.environ.setdefault("TORCHINDUCTOR_DISABLED", "1")
 
-# ── remote‑code flag ───────────────────────────────────────────────────
+# ───────────────────────  remote‑code flag  ───────────────────────────
 os.environ.setdefault("HF_ALLOW_CODE_IMPORT", "1")
 TOKEN_KW = {"trust_remote_code": True}
 
-# ── config ─────────────────────────────────────────────────────────────
+# ───────────────────────────  config  ─────────────────────────────────
 DEVICE      = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 WEIGHT_REPO = "Sleepyriizi/Orify-Text-Detection-Weights"
 FILE_MAP    = {"ensamble_1":"ensamble_1","ensamble_2.bin":"ensamble_2.bin","ensamble_3":"ensamble_3"}
@@ -38,7 +38,7 @@ LABELS      = {i:n for i,n in enumerate([
     "mixtral-8x7b","opt-1.3b","opt-125m","opt-13b","opt-2.7b","opt-30b","opt-350m","opt-6.7b",
     "opt-iml-30b","opt-iml-max-1.3b","t0-11b","t0-3b","text-davinci-002","text-davinci-003"])}
 
-# ── JWT helpers ───────────────────────────────────────────────────────
+# ────────────────────────  JWT helpers  ──────────────────────────────
 SECRET_KEY = os.getenv("SECRET_KEY")
 if not SECRET_KEY:
     raise RuntimeError("SECRET_KEY env‑var not set – add it in Space settings → Secrets")
@@ -55,32 +55,30 @@ def _verify_jwt(tok:str=Depends(oauth2)):
     except JWTError:
         raise HTTPException(401,"Invalid or expired token")
 
-# ── model bootstrap ───────────────────────────────────────────────────
+# ───────────────────────  model bootstrap  ───────────────────────────
 print("🔄 Fetching ensemble weights…", flush=True)
 paths={k:hf_hub_download(WEIGHT_REPO,f,resume_download=True) for k,f in FILE_MAP.items()}
 
 print("🧩 Building ModernBERT backbone…", flush=True)
-_cfg = AutoConfig.from_pretrained(BASE_MODEL, **TOKEN_KW)
-_cfg.num_labels = NUM_LABELS                       # ➜ classification head = 41
+_cfg = AutoConfig.from_pretrained(BASE_MODEL, **TOKEN_KW); _cfg.num_labels = NUM_LABELS
 _tok = AutoTokenizer.from_pretrained(BASE_MODEL, **TOKEN_KW)
 _models: List[AutoModelForSequenceClassification] = []
 for p in paths.values():
     m = AutoModelForSequenceClassification.from_pretrained(
             BASE_MODEL,
             config=_cfg,
-            ignore_mismatched_sizes=True,          # skip the 2‑class head in checkpoint
+            ignore_mismatched_sizes=True,
             **TOKEN_KW,
         )
-    state=torch.load(p, map_location=DEVICE)
-    m.load_state_dict(state)                       # loads 41‑class ensemble head
+    m.load_state_dict(torch.load(p, map_location=DEVICE))
     m.to(DEVICE).eval()
     _models.append(m)
 print(f"✅ Ensemble ready on {DEVICE}")
 
-# ── helper fns ─────────────────────────────────────────────────────────
+# ─────────────────────────  helpers  ─────────────────────────────────
 
 def _tidy(t:str)->str:
-    t=t.replace("\r\n","\n").replace("\r","\n")
+    t=t.replace("\r\n","\n").replace("\r", "\n")
     t=re.sub(r"\n\s*\n+","\n\n",t)
     t=re.sub(r"[ \t]+"," ",t)
     t=re.sub(r"(\w+)-\n(\w+)",r"\1\2",t)
@@ -96,14 +94,14 @@ def _infer(seg:str):
     top3=[LABELS[i] for i in torch.topk(ai_probs,3).indices.tolist()]
     return human, ai, top3
 
-# ── schemas ────────────────────────────────────────���──────────────────
+# ─────────────────────────  schemas  ─────────────────────────────────
 class TokenOut(BaseModel): access_token:str; token_type:str="bearer"
 class AnalyseIn(BaseModel): text:str=Field(...,min_length=1)
 class Line(BaseModel): text:str; ai:float; human:float; top3:List[str]; reason:str
 class AnalyseOut(BaseModel): verdict:str; confidence:float; ai_avg:float; human_avg:float; per_line:List[Line]; highlight_html:str
 
-# ── FastAPI app ───────────────────────────────────────────────────────
-app=FastAPI(title="Orify Text Detector API",version="1.1.1")
+# ─────────────────────────  FastAPI app  ─────────────────────────────
+app=FastAPI(title="Orify Text Detector API",version="1.2.0")
 app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_methods=["*"],allow_headers=["*"])
 
 @app.post("/token",response_model=TokenOut)
@@ -127,3 +125,9 @@ async def analyse(body:AnalyseIn,_=Depends(_verify_jwt)):
     badge=(f"<span class='ai-line' style='padding:6px 10px;font-weight:bold'>AI-generated {ai_avg:.2f}%</span>" if verdict=="AI-generated" else f"<span class='human-line' style='padding:6px 10px;font-weight:bold'>Human-written {human_avg:.2f}%</span>")
     html_out=f"<h3>{badge}</h3><hr>"+"<br>".join(html_parts)
     return AnalyseOut(verdict=verdict,confidence=conf,ai_avg=ai_avg,human_avg=human_avg,per_line=per,highlight_html=html_out)
+
+# ───────────────────────────  entrypoint  ────────────────────────────
+if __name__ == "__main__":
+    import uvicorn, sys
+    port=int(os.environ.get("PORT", "7860"))
+    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="info", reload=False)