Update app.py

app.py

@@ -6,7 +6,7 @@ for detecting and visualizing PII in PDF/DOC/DOCX documents.
 
 - Backend: gr.Server (Gradio + FastAPI)
 - Frontend: Custom HTML/CSS/JS
-- Model: openai/openai-privacy-filter (1.5B params, 50M active, 128k context)
+- Model: charles-first-org/second-model (1.5B params, 50M active, 128k context)
 """
 
 import os
@@ -21,7 +21,8 @@ from fastapi import UploadFile, File
 from fastapi.responses import HTMLResponse, JSONResponse
 
 # ── Configuration ────────────────────────────────────────────────
-MODEL_ID = os.getenv("MODEL_ID", "charles-first-org/second-model") #"openai/openai-privacy-filter")
+MODEL_ID = os.getenv("MODEL_ID", "charles-first-org/second-model")
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 CATEGORIES = {
     "private_person":  {"color": "#ef4444", "bg": "rgba(239,68,68,0.15)",  "label": "Person"},
@@ -40,9 +41,10 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from transformers import AutoTokenizer, AutoModelForTokenClassification  # noqa: E402
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
 model = AutoModelForTokenClassification.from_pretrained(
-    MODEL_ID, trust_remote_code=True, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
+    MODEL_ID, trust_remote_code=True, token=HF_TOKEN,
+    torch_dtype=torch.bfloat16 if device.type == "cuda" else torch.float32,
 )
 model.eval().to(device)