Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import HTMLResponse | |
| from pydantic import BaseModel | |
| from transformers import pipeline | |
| import uvicorn | |
| app = FastAPI( | |
| title="PII Warden Cloud AI Endpoint", | |
| description="Tier 2 Cloud AI Inference service for the PII Warden browser extension." | |
| ) | |
| # Enable CORS (Cross-Origin Resource Sharing) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| class AnalyzeRequest(BaseModel): | |
| text: str | |
| # Load the standard PyTorch model from Hugging Face on startup | |
| print("Loading DistilBERT PII model into memory...") | |
| try: | |
| nlp_pipeline = pipeline( | |
| "token-classification", | |
| model="samuelolubukun/pii-ner-finetuned-distilbert", | |
| aggregation_strategy="simple" | |
| ) | |
| print("Model loaded successfully!") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| nlp_pipeline = None | |
| async def read_root(): | |
| return """ | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>PII Warden AI Endpoint</title> | |
| <link href="https://fonts.googleapis.com/css2?family=Outfit:wght@400;600;700&display=swap" rel="stylesheet"> | |
| <style> | |
| body { | |
| font-family: 'Outfit', sans-serif; | |
| background-color: #0b0f19; | |
| color: #f3f4f6; | |
| display: flex; | |
| flex-direction: column; | |
| justify-content: center; | |
| align-items: center; | |
| min-height: 100vh; | |
| margin: 0; | |
| background: radial-gradient(circle at top right, rgba(139, 92, 246, 0.15), transparent 60%), | |
| radial-gradient(circle at bottom left, rgba(236, 72, 153, 0.1), transparent 60%); | |
| } | |
| .card { | |
| background: rgba(17, 24, 39, 0.75); | |
| border: 1px solid rgba(255, 255, 255, 0.08); | |
| border-radius: 16px; | |
| padding: 32px; | |
| max-width: 500px; | |
| text-align: center; | |
| backdrop-filter: blur(12px); | |
| box-shadow: 0 8px 32px rgba(0, 0, 0, 0.5); | |
| } | |
| h1 { | |
| font-size: 2rem; | |
| margin-bottom: 8px; | |
| background: linear-gradient(135deg, #ffffff 30%, #a78bfa 100%); | |
| -webkit-background-clip: text; | |
| -webkit-text-fill-color: transparent; | |
| } | |
| p { | |
| color: #9ca3af; | |
| font-size: 0.95rem; | |
| line-height: 1.5; | |
| } | |
| .badge { | |
| display: inline-block; | |
| background: rgba(16, 185, 129, 0.15); | |
| color: #10b981; | |
| border: 1px solid rgba(16, 185, 129, 0.3); | |
| padding: 6px 16px; | |
| border-radius: 9999px; | |
| font-size: 0.8rem; | |
| font-weight: 600; | |
| text-transform: uppercase; | |
| letter-spacing: 0.05em; | |
| margin-top: 16px; | |
| } | |
| .endpoint { | |
| background: rgba(10, 10, 10, 0.4); | |
| padding: 10px; | |
| border-radius: 8px; | |
| font-family: monospace; | |
| font-size: 0.85rem; | |
| color: #f472b6; | |
| margin-top: 20px; | |
| border: 1px solid rgba(255, 255, 255, 0.05); | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="card"> | |
| <h1>🛡️ PII Warden AI</h1> | |
| <p>Your hosted PII inference endpoint is live. Configure your browser extension to query the endpoint below for Tier 2 context analysis.</p> | |
| <div class="endpoint">POST /analyze</div> | |
| <div class="badge">Online & Active</div> | |
| </div> | |
| </body> | |
| </html> | |
| """ | |
| async def analyze_text(request: AnalyzeRequest): | |
| if nlp_pipeline is None: | |
| raise HTTPException(status_code=503, detail="Model pipeline not initialized.") | |
| try: | |
| text = request.text | |
| if not text.strip(): | |
| return [] | |
| predictions = nlp_pipeline(text) | |
| formatted_results = [] | |
| for pred in predictions: | |
| formatted_results.append({ | |
| "entity_group": pred["entity_group"], | |
| "score": float(pred["score"]), | |
| "word": pred["word"], | |
| "start": int(pred["start"]), | |
| "end": int(pred["end"]) | |
| }) | |
| return formatted_results | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}") | |
| async def health_check(): | |
| return {"status": "healthy", "model_loaded": nlp_pipeline is not None} | |
| if __name__ == "__main__": | |
| # Hugging Face Spaces require listening on port 7860 | |
| uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False) | |