Update app.py

app.py

@@ -1,10 +1,15 @@
+
 from __future__ import annotations
 import os, re, html
 from datetime import datetime, timedelta
 from typing import List
 
 import torch
-from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+)
 from huggingface_hub import hf_hub_download
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
@@ -12,23 +17,22 @@ from fastapi.middleware.cors import CORSMiddleware
 from jose import jwt, JWTError
 from pydantic import BaseModel, Field
 
-# ── torch.compile shim (optional) ──────────────────────────────────────
+# ── torch shim ─────────────────────────────────────────────────────────
 if hasattr(torch, "compile"):
     torch.compile = (lambda m=None,*_,**__: m if callable(m) else (lambda f: f))  # type: ignore
-    # keep Inductor disabled on CPU Spaces; harmless on GPU
     os.environ.setdefault("TORCHINDUCTOR_DISABLED", "1")
 
-# ── allow ModernBERT remote code ──────────────────────────────────────
-os.environ.setdefault("HF_ALLOW_CODE_IMPORT", "1")  # required for custom architecture
+# ── 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"}
+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([
+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",
@@ -38,81 +42,89 @@ LABELS = {i:n for i,n in enumerate([
 # ── JWT helpers ───────────────────────────────────────────────────────
 SECRET_KEY = os.getenv("SECRET_KEY")
 if not SECRET_KEY:
-    raise RuntimeError("SECRET_KEY env‑var not set (add it in the Space ➜ Settings ➜ Secrets)")
-ALG = "HS256"; EXP_HOURS = 24
+    raise RuntimeError("SECRET_KEY env‑var not set – add it in Space settings → Secrets")
+ALG="HS256"; EXP=24
 oauth2 = OAuth2PasswordBearer(tokenUrl="token")
 
-def _make_token(sub:str)->str:
-    return jwt.encode({"sub":sub,"exp":datetime.utcnow()+timedelta(hours=EXP_HOURS)}, SECRET_KEY, algorithm=ALG)
+def _make_jwt(sub:str)->str:
+    payload={"sub":sub,"exp":datetime.utcnow()+timedelta(hours=EXP)}
+    return jwt.encode(payload,SECRET_KEY,algorithm=ALG)
 
-def _verify(tok:str=Depends(oauth2)):
+def _verify_jwt(tok:str=Depends(oauth2)):
     try:
-        return jwt.decode(tok, SECRET_KEY, algorithms=[ALG])["sub"]
+        return jwt.decode(tok,SECRET_KEY,algorithms=[ALG])["sub"]
     except JWTError:
-        raise HTTPException(401, "Invalid or expired token")
+        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()}
 
-# ── load ensemble ─────────────────────────────────────────────────────
-print("🔄 Downloading weights…", flush=True)
-local_paths = {k: hf_hub_download(WEIGHT_REPO, f, resume_download=True) for k, f in FILE_MAP.items()}
-print("🧩 Loading ModernBERT checkpoints…", flush=True)
+print("🧩 Building ModernBERT backbone…", flush=True)
 _cfg = AutoConfig.from_pretrained(BASE_MODEL, **TOKEN_KW)
+_cfg.num_labels = NUM_LABELS                       # ➜ classification head = 41
 _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))
+for p in paths.values():
+    m = AutoModelForSequenceClassification.from_pretrained(
+            BASE_MODEL,
+            config=_cfg,
+            ignore_mismatched_sizes=True,          # skip the 2‑class head in checkpoint
+            **TOKEN_KW,
+        )
+    state=torch.load(p, map_location=DEVICE)
+    m.load_state_dict(state)                       # loads 41‑class ensemble head
     m.to(DEVICE).eval()
     _models.append(m)
 print(f"✅ Ensemble ready on {DEVICE}")
 
-# ── helpers ───────────────────────────────────────────────────────────
+# ── helper fns ─────────────────────────────────────────────────────────
 
 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)
+    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)
+    inp=_tok(seg,return_tensors="pt",truncation=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]
+        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()]
+    top3=[LABELS[i] for i in torch.topk(ai_probs,3).indices.tolist()]
     return human, ai, top3
 
-# ── Pydantic schemas ─────────────────────────────────────────────────
+# ── schemas ───────────────────────────────────────────────────────────
 class TokenOut(BaseModel): access_token:str; token_type:str="bearer"
-class AnalyseIn(BaseModel): text:str=Field(..., min_length=1)
+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.0")
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+app=FastAPI(title="Orify Text Detector API",version="1.1.1")
+app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_methods=["*"],allow_headers=["*"])
 
-@app.post("/token", response_model=TokenOut)
-async def token_route(form: OAuth2PasswordRequestForm = Depends()):
-    return TokenOut(access_token=_make_token(form.username))
+@app.post("/token",response_model=TokenOut)
+async def token(form:OAuth2PasswordRequestForm=Depends()):
+    return TokenOut(access_token=_make_jwt(form.username))
 
-@app.post("/analyse", response_model=AnalyseOut)
-async def analyse(req: AnalyseIn, _user=Depends(_verify)):
-    lines=_tidy(req.text).split("\n"); html_parts=[]; records=[]; h_sum=ai_sum=n=0.0
+@app.post("/analyse",response_model=AnalyseOut)
+async def analyse(body:AnalyseIn,_=Depends(_verify_jwt)):
+    lines=_tidy(body.text).split("\n"); html_parts=[]; per=[]; 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}%"
+        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}%)")
-        records.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"; conf=max(human_avg, ai_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>")
-    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=records, highlight_html=html_out)
-
+        reason=(f"High AI likelihood ({ai:.1f}%) – fingerprint ≈ {top3[0]}" if ai>human else f"Lexical variety suggests human ({human:.1f}%)")
+        per.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"; conf=max(human_avg,ai_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>")
+    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)