Spaces:

samuelolubukun
/

PII-Warden-API

Sleeping

App Files Files Community

samuelolubukun commited on Jun 17

Commit

cecc83f

verified ·

1 Parent(s): e3f98df

Upload 4 files

Browse files

Files changed (3) hide show

Dockerfile +2 -2
app.py +10 -18
requirements.txt +0 -2

Dockerfile CHANGED Viewed

@@ -19,9 +19,9 @@ COPY --chown=user requirements.txt .
 # Install dependencies using CPU-only extra index to prevent massive GPU CUDA installations
 RUN pip install --no-cache-dir --upgrade -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
-# Pre-cache the ONNX model during the image build process
 # This guarantees instant startup times when the Space boots up.
-RUN python -c "from optimum.onnxruntime import ORTModelForTokenClassification; ORTModelForTokenClassification.from_pretrained('samuelolubukun/pii-ner-edge-optimized', file_name='model_quantized.onnx')"
 # Copy the rest of the application files
 COPY --chown=user . .

 # Install dependencies using CPU-only extra index to prevent massive GPU CUDA installations
 RUN pip install --no-cache-dir --upgrade -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cpu
+# Pre-cache the standard PyTorch DistilBERT model during the image build process
 # This guarantees instant startup times when the Space boots up.
+RUN python -c "from transformers import pipeline; pipeline('token-classification', model='samuelolubukun/pii-ner-finetuned-distilbert')"
 # Copy the rest of the application files
 COPY --chown=user . .

app.py CHANGED Viewed

@@ -2,8 +2,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
-from transformers import AutoTokenizer, pipeline
-from optimum.onnxruntime import ORTModelForTokenClassification
 import uvicorn
 app = FastAPI(
@@ -23,24 +22,17 @@ app.add_middleware(
 class AnalyzeRequest(BaseModel):
     text: str
-# Load the optimized ONNX model from Hugging Face on startup
-print("Loading ONNX DistilBERT PII model into memory...")
 try:
-    model = ORTModelForTokenClassification.from_pretrained(
-        "samuelolubukun/pii-ner-edge-optimized",
-        file_name="model_quantized.onnx"
-    )
-    tokenizer = AutoTokenizer.from_pretrained("samuelolubukun/pii-ner-edge-optimized")
     nlp_pipeline = pipeline(
         "token-classification",
-        model=model,
-        tokenizer=tokenizer,
         aggregation_strategy="simple"
     )
-    print("ONNX Model loaded successfully!")
 except Exception as e:
-    print(f"Error loading ONNX model: {e}")
     nlp_pipeline = None
 @app.get("/", response_class=HTMLResponse)
@@ -114,8 +106,8 @@ async def read_root():
     </head>
     <body>
         <div class="card">
-            <h1>🛡️ PII Warden AI (ONNX)</h1>
-            <p>Your hosted edge-quantized ONNX 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>
@@ -126,7 +118,7 @@ async def read_root():
 @app.post("/analyze")
 async def analyze_text(request: AnalyzeRequest):
     if nlp_pipeline is None:
-        raise HTTPException(status_code=503, detail="ONNX Model pipeline not initialized.")
     try:
         text = request.text
@@ -148,7 +140,7 @@ async def analyze_text(request: AnalyzeRequest):
         return formatted_results
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"ONNX Inference error: {str(e)}")
 @app.get("/health")
 async def health_check():

 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
+from transformers import pipeline
 import uvicorn
 app = FastAPI(
 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)
     </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>
 @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
         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():

requirements.txt CHANGED Viewed

@@ -2,7 +2,5 @@ fastapi==0.110.0
 uvicorn==0.29.0
 pydantic==2.6.4
 transformers==4.38.2
-optimum==1.18.0
-onnxruntime==1.17.1
 numpy<2
 torch==2.2.1

 uvicorn==0.29.0
 pydantic==2.6.4
 transformers==4.38.2
 numpy<2
 torch==2.2.1