Spaces:
Sleeping
Sleeping
File size: 3,497 Bytes
be6176b 2f44304 be6176b 8b49703 be6176b 9b47fbd be6176b 2f44304 be6176b | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from gradio_client import Client
import re
import time
app = FastAPI(title="PrivBERT Proxy")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
PRIVBERT_SPACE = "shikhasoneji8/privacywhisper"
class AnalyzeRequest(BaseModel):
text: str
threshold: float = 0.75
def extract_choices(result_value) -> list[str]:
"""Extract topic strings from a Gradio Dropdown result (handles dict or list)."""
if isinstance(result_value, dict):
choices = result_value.get("choices", [])
return [c[0] if isinstance(c, list) else c for c in choices if c and (c[0] if isinstance(c, list) else c) != "—"]
if isinstance(result_value, list):
return [c for c in result_value if c and c != "—"]
return []
def count_clauses(html: str) -> int:
return len(re.findall(r"class='card ", html))
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/analyze")
def analyze(req: AnalyzeRequest):
if not req.text or len(req.text.strip()) < 50:
raise HTTPException(status_code=400, detail="Text too short — provide at least 50 characters of policy text")
# Each request gets its own client so sessions don't bleed across requests.
# Retry up to 3 times — HF free-tier spaces can take a moment to wake up.
client = None
last_error = ""
for attempt in range(3):
try:
client = Client(PRIVBERT_SPACE, verbose=False)
break
except Exception as e:
last_error = str(e)
if attempt < 2:
time.sleep(5)
if client is None:
raise HTTPException(status_code=503, detail=f"Could not connect to PrivBERT model after 3 attempts: {last_error}")
# Step 1: Run annotation — classifies text into privacy topics
try:
result1 = client.predict(
text_input=req.text,
threshold=req.threshold,
api_name="/_run_annotation",
)
except Exception as e:
raise HTTPException(status_code=502, detail=f"_run_annotation failed: {str(e)}")
# result1 is a tuple: (topics_result, granular_result, markdown_placeholder)
topics = extract_choices(result1[0]) if len(result1) > 0 else []
granular_topics = extract_choices(result1[1]) if len(result1) > 1 else []
# Step 2: Render annotated HTML — same client = same session = server has state from step 1
annotated_html = ""
try:
result2 = client.predict(
selected_topics=topics,
selected_fine=granular_topics,
selected_ratings=[],
selected_compliance=[],
view_mode="Cards",
fsize=15,
enable_expl=False,
api_name="/_apply_filters_and_render",
)
annotated_html = result2[0] if result2 and isinstance(result2[0], str) else ""
except Exception as e:
# Annotated HTML is best-effort — still return topics if this fails
print(f"Warning: _apply_filters_and_render failed: {e}")
return {
"success": True,
"topics": topics,
"granular_topics": granular_topics,
"annotated_html": annotated_html,
"clause_count": count_clauses(annotated_html),
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|