Spaces:
Sleeping
Sleeping
File size: 3,388 Bytes
a6b677f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import pathlib, pickle, pdfplumber, spacy, torch, tempfile, re, io
app = FastAPI(title="Innocence-Claim API", version="1.0")
# ---------- CORS Configuration ----------
# Configured for Hugging Face Spaces - allows all origins for API accessibility
# Hugging Face Spaces URLs follow pattern: https://{username}-{space-name}.hf.space
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Allow all origins for public API access
allow_credentials=False, # Must be False when allow_origins is ["*"]
allow_methods=["GET", "POST", "OPTIONS"],
allow_headers=["*"],
)
# ---------- load pipeline ----------
pkl_path = pathlib.Path(__file__).parent.parent / "models" / "innocence_pipeline.pkl"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Custom unpickler to handle CUDA tensors on CPU
class CPU_Unpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'torch.storage' and name == '_load_from_bytes':
return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
else:
return super().find_class(module, name)
with open(pkl_path, "rb") as f:
if device.type == "cpu":
bundle = CPU_Unpickler(f).load()
else:
bundle = pickle.load(f)
tokenizer, model = bundle["tokenizer"], bundle["model"]
model.to(device)
nlp = spacy.load("en_core_web_sm")
def predict(text: str) -> float:
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
with torch.no_grad():
prob = torch.softmax(model(**inputs).logits, dim=1)[0, 1].item()
return round(prob, 3)
def analyse_pdf(pdf_path: pathlib.Path, cutoff: float):
rows, total, score_sum = [], 0, 0.0
with pdfplumber.open(pdf_path) as pdf:
for page_idx, page in enumerate(pdf.pages, 1):
text = page.extract_text() or ""
for sent_idx, sent in enumerate(nlp(text).sents, 1):
s = sent.text.strip()
if 10 < len(s) < 500:
total += 1
score = predict(s)
score_sum += score
if score >= cutoff:
rows.append({
"sentence": s,
"confidence": score,
"page": page_idx,
"sent_id": sent_idx,
})
reliability = round((score_sum / total) * 100, 1) if total else 0.0
tier = "High" if reliability >= 80 else "Medium" if reliability >= 50 else "Low"
return {"reliability_percent": reliability, "tier": tier}
@app.post("/predict")
async def predict_pdf(
file: UploadFile = File(...),
cutoff: float = Form(0.7),
):
if file.content_type != "application/pdf":
raise HTTPException(400, "PDF required")
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(await file.read())
tmp_path = pathlib.Path(tmp.name)
try:
return analyse_pdf(tmp_path, cutoff)
finally:
tmp_path.unlink(missing_ok=True)
@app.get("/health")
def health():
return {"status": "ok"} |