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 @app.get("/", response_class=HTMLResponse) async def read_root(): return """ PII Warden AI Endpoint

🛡️ PII Warden AI

Your hosted PII inference endpoint is live. Configure your browser extension to query the endpoint below for Tier 2 context analysis.

POST /analyze

Online & Active

""" @app.post("/analyze") 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)}") @app.get("/health") 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)