Update app.py

app.py

@@ -1,3 +1,5 @@
+"""Orify Text Detector API – FastAPI / JWT / HF Zero-GPU"""
+
 from __future__ import annotations
 import os, re, html
 from datetime import datetime, timedelta
@@ -5,169 +7,126 @@ from typing import List
 
 import torch
 from transformers import (
-    AutoTokenizer,
     AutoConfig,
+    AutoTokenizer,
     AutoModelForSequenceClassification,
 )
 from huggingface_hub import hf_hub_download
-from fastapi import FastAPI, HTTPException, Depends
+from fastapi import FastAPI, Depends, HTTPException
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from fastapi.middleware.cors import CORSMiddleware
 from jose import jwt, JWTError
 from pydantic import BaseModel, Field
 
-# ---------------------------------------------------------------------
-# Torch compile shim (CPU runtime)
-# ---------------------------------------------------------------------
-if hasattr(torch, "compile"):
-    torch.compile = (lambda m=None, *_, **__: m if callable(m) else (lambda f: f))  # type: ignore
+# ─────────────────────  Torch shim  ─────────────────────
+if hasattr(torch, "compile"):                        # HF CPU image has torch-compile
+    torch.compile = lambda m=None,*_,**__: m if callable(m) else (lambda f: f)  # type: ignore
     os.environ["TORCHINDUCTOR_DISABLED"] = "1"
 
-# ---------------------------------------------------------------------
-# Environment flags – enable remote code in Transformers
-# ---------------------------------------------------------------------
-os.environ.setdefault("HF_ALLOW_CODE_IMPORT", "1")  # allow custom ModernBERT classes
-
-# ---------------------------------------------------------------------
-# Model / weight 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"}
-BASE_MODEL  = "answerdotai/ModernBERT-base"
-NUM_LABELS  = 41
-TOKEN_KW    = {"trust_remote_code": True}
-
-LABELS = {i: n for i, n in enumerate([
-    "13B", "30B", "65B", "7B", "GLM130B", "bloom_7b", "bloomz", "cohere", "davinci",
-    "dolly", "dolly-v2-12b", "flan_t5_base", "flan_t5_large", "flan_t5_small",
-    "flan_t5_xl", "flan_t5_xxl", "gemma-7b-it", "gemma2-9b-it", "gpt-3.5-turbo",
-    "gpt-35", "gpt-4", "gpt-4o", "gpt-j", "gpt-neox", "human", "llama3-70b",
-    "llama3-8b", "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",
+# ─────────────────────  Remote-code flag  ───────────────
+os.environ.setdefault("HF_ALLOW_CODE_IMPORT", "1")   # let Transformers load ModernBERT code
+TOKEN_KW = {"trust_remote_code": True}
+
+# ─────────────────────  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"}
+BASE_MODEL   = "answerdotai/ModernBERT-base"
+NUM_LABELS   = 41
+
+LABELS = {i:n for i,n in enumerate([
+    "13B","30B","65B","7B","GLM130B","bloom_7b","bloomz","cohere","davinci","dolly","dolly-v2-12b",
+    "flan_t5_base","flan_t5_large","flan_t5_small","flan_t5_xl","flan_t5_xxl","gemma-7b-it","gemma2-9b-it",
+    "gpt-3.5-turbo","gpt-35","gpt-4","gpt-4o","gpt-j","gpt-neox","human","llama3-70b","llama3-8b",
+    "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
-# ---------------------------------------------------------------------
-SECRET_KEY = os.getenv("SECRET_KEY")
+# ─────────────────────  JWT helpers  ────────────────────
+SECRET_KEY = os.getenv("SECRET_KEY")           # set this in Space → Settings → Secrets
 if not SECRET_KEY:
-    raise RuntimeError("SECRET_KEY env‑var not set")
-ALGO = "HS256"
-EXP_H = 24
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+    raise RuntimeError("Set SECRET_KEY env-var")
+ALG          = "HS256"
+EXP_HOURS    = 24
+oauth_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
-def _jwt_create(sub: str) -> str:
-    return jwt.encode({"sub": sub, "exp": datetime.utcnow() + timedelta(hours=EXP_H)}, SECRET_KEY, algorithm=ALGO)
+def _jwt_make(sub:str)->str:
+    payload={"sub":sub,"exp":datetime.utcnow()+timedelta(hours=EXP_HOURS)}
+    return jwt.encode(payload, SECRET_KEY, algorithm=ALG)
 
-def _jwt_verify(tok: str = Depends(oauth2_scheme)) -> str:
+def _jwt_ok(token:str=Depends(oauth_scheme))->str:
     try:
-        return jwt.decode(tok, SECRET_KEY, algorithms=[ALGO])["sub"]
+        return jwt.decode(token, SECRET_KEY, algorithms=[ALG])["sub"]
     except JWTError:
-        raise HTTPException(401, "Invalid or expired token")
-
-# ---------------------------------------------------------------------
-# Load tokenizer + config + ensemble
-# ---------------------------------------------------------------------
-print("🔄 Downloading weights…", flush=True)
-local_paths = {k: hf_hub_download(WEIGHT_REPO, v, resume_download=True) for k, v in FILE_MAP.items()}
-
-print("🧩 Loading ModernBERT remote code…", flush=True)
-_cfg = AutoConfig.from_pretrained(BASE_MODEL, **TOKEN_KW)
-_tok = AutoTokenizer.from_pretrained(BASE_MODEL, **TOKEN_KW)
-_models: List[AutoModelForSequenceClassification] = []
-for p in local_paths.values():
-    m = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL, config=_cfg, **TOKEN_KW)
-    m.load_state_dict(torch.load(p, map_location=DEVICE))
-    m.to(DEVICE).eval()
-    _models.append(m)
+        raise HTTPException(401,"Invalid or expired token")
+
+# ─────────────────────  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("🧩 Loading ModernBERT code…", flush=True)
+cfg  = AutoConfig.from_pretrained(BASE_MODEL, **TOKEN_KW)
+tok  = AutoTokenizer.from_pretrained(BASE_MODEL, **TOKEN_KW)
+models=[]
+for p in paths.values():
+    m = AutoModelForSequenceClassification.from_pretrained(BASE_MODEL, config=cfg, **TOKEN_KW)
+    m.load_state_dict(torch.load(p,map_location=DEVICE))
+    m.to(DEVICE).eval(); models.append(m)
 print("✅ Ensemble ready")
 
-# ---------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------
-
-def _tidy(t: str) -> str:
-    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)
-    t = re.sub(r"(?<!\n)\n(?!\n)", " ", t)
+# ─────────────────────  Utils  ──────────────────────────
+def _tidy(t:str)->str:
+    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)
+    t=re.sub(r"(?<!\n)\n(?!\n)"," ",t)
     return t.strip()
 
-def _infer(seg: str):
-    inp = _tok(seg, return_tensors="pt", truncation=True, padding=True).to(DEVICE)
+def _infer(seg:str):
+    inp=tok(seg,return_tensors="pt",trunc=True,padding=True).to(DEVICE)
     with torch.no_grad():
-        probs = torch.stack([torch.softmax(m(**inp).logits, dim=1) for m in _models]).mean(0)[0]
-    ai_probs = probs.clone(); ai_probs[24] = 0
-    ai = ai_probs.sum().item() * 100
-    human = 100 - ai
-    top3 = [LABELS[i] for i in torch.topk(ai_probs, 3).indices.tolist()]
-    return human, ai, top3
-
-# ---------------------------------------------------------------------
-# Schemas
-# ---------------------------------------------------------------------
-class Token(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
-
+        probs=torch.stack([torch.softmax(m(**inp).logits,1) for m in models]).mean(0)[0]
+    ai_probs=probs.clone(); ai_probs[24]=0
+    ai=ai_probs.sum().item()*100; human=100-ai
+    top3=[LABELS[i] for i in torch.topk(ai_probs,3).indices.tolist()]
+    return human,ai,top3
+
+# ─────────────────────  Schemas  ────────────────────────
+class Token(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 = FastAPI(title="Orify Text Detector API", version="1.0.0")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
-
-@app.post("/token", response_model=Token)
-async def login(form: OAuth2PasswordRequestForm = Depends()):
-    return Token(access_token=_jwt_create(form.username))
-
-@app.post("/analyse", response_model=AnalyseOut)
-async def analyse(data: AnalyseIn, _=Depends(_jwt_verify)):
-    lines = _tidy(data.text).split("\n")
-    per_line, html_parts = [], []
-    h_sum = ai_sum = n = 0.0
+    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.0.0")
+app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_methods=["*"],allow_headers=["*"])
+
+@app.post("/token",response_model=Token,summary="Get JWT")
+async def login(form:OAuth2PasswordRequestForm=Depends()):
+    return Token(access_token=_jwt_make(form.username))
 
+@app.post("/analyse",response_model=AnalyseOut,summary="Detect AI-generated text")
+async def analyse(req:AnalyseIn,_user=Depends(_jwt_ok)):
+    lines=_tidy(req.text).split("\\n")
+    html_parts=[]; per_line=[]; h_sum=ai_sum=n=0.0
     for ln in lines:
-        if not ln.strip():
-            html_parts.append("<br>")
-            continue
-        n += 1
-        human, ai, top3 = _infer(ln)
-        h_sum += human; ai_sum += ai
-        cls = "ai-line" if ai > human else "human-line"
-        tip = f"AI {ai:.2f}% – Top-3: {', '.join(top3)}" if ai > human else f"Human {human:.2f}%"
-        html_parts.append(f"<span class='{cls} prob-tooltip' title='{tip}'>{html.escape(ln)}</span>")
-        reason = (f"High AI likelihood ({ai:.1f}%) – fingerprint ≈ {top3[0]}" if ai > human else
-                  f"Lexical variety suggests human ({human:.1f}%)")
-        per_line.append(Line(text=ln, ai=ai, human=human, top3=top3, reason=reason))
-
-    human_avg = h_sum / n if n else 0
-    ai_avg    = ai_sum / n if n else 0
-    verdict   = "AI-generated" if ai_avg > human_avg else "Human-written"
-    confidence = max(ai_avg, human_avg)
-    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>")
-    highlight_html = f"<h3>{badge}</h3><hr>" + "<br>".join(html_parts)
-
-    return AnalyseOut(verdict=verdict, confidence=confidence, ai_avg=ai_avg, human_avg=human_avg,
-                      per_line=per_line, highlight_html=highlight_html)
+        if not ln.strip(): html_parts.append("<br>"); continue
+        n+=1
+        h,ai,top3=_infer(ln); h_sum+=h; ai_sum+=ai
+        cls="ai-line" if ai>h else "human-line"
+        tip=f\"AI {ai:.2f}% – Top-3: {', '.join(top3)}\" if ai>h else f\"Human {h:.2f}%\"
+        html_parts.append(f\"<span class='{cls} prob-tooltip' title='{tip}'>{html.escape(ln)}</span>\")
+        reason=(f\"High AI likelihood ({ai:.1f}%) – fingerprint ≈ {top3[0]}\" if ai>h
+                else f\"Lexical variety suggests human ({h:.1f}%)\")
+        per_line.append(Line(text=ln,ai=ai,human=h,top3=top3,reason=reason))
+    human_avg=h_sum/n if n else 0; ai_avg=ai_sum/n if n else 0
+    verdict=\"AI-generated\" if ai_avg>human_avg else \"Human-written\"
+    conf=max(ai_avg,human_avg)
+    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>\")
+    return AnalyseOut(verdict=verdict,confidence=conf,ai_avg=ai_avg,human_avg=human_avg,
+                      per_line=per_line,highlight_html=f\"<h3>{badge}</h3><hr>\"+\"<br>\".join(html_parts))