Update app.py

app.py

@@ -10,6 +10,7 @@ from PyPDF2 import PdfReader, PdfWriter
 from presidio_analyzer import AnalyzerEngine, RecognizerRegistry, PatternRecognizer, RecognizerResult
 from presidio_anonymizer import AnonymizerEngine
 from presidio_anonymizer.entities import OperatorConfig
+import tempfile
 
 st.set_page_config(page_title="Presidio PHI De-identification", layout="wide", initial_sidebar_state="expanded", menu_items={"About": "https://microsoft.github.io/presidio/"})
 dotenv.load_dotenv()
@@ -44,14 +45,10 @@ def analyzer_engine(model_family: str, model_path: str) -> AnalyzerEngine:
 
 def get_supported_entities(model_family: str, model_path: str) -> list[str]:
     """📋 Lists what secrets we’re hunting—PHI beware!"""
-    if model_family.lower() == "huggingface":
-        return ["PERSON", "LOCATION", "ORGANIZATION", "DATE_TIME"]
-    elif model_family.lower() == "flair":
-        return ["PERSON", "LOCATION", "ORGANIZATION"]
-    return ["PERSON", "LOCATION", "ORGANIZATION"]
+    return ["PERSON", "LOCATION", "ORGANIZATION", "DATE_TIME"] if model_family.lower() == "huggingface" else ["PERSON", "LOCATION", "ORGANIZATION"]
 
 # Feature Spotlight: 🕵️‍♂️ The Great PHI Hunt Begins!
-# With a flick of the wrist, we summon models to sniff out sensitive data in PDFs, making privacy a breeze! 😎
+# Summon models to sniff out sensitive data in PDFs with ninja stealth! 😎
 
 def analyze(analyzer: AnalyzerEngine, text: str, entities: list[str], language: str, score_threshold: float, return_decision_process: bool, allow_list: list[str], deny_list: list[str]) -> list[RecognizerResult]:
     """🦸 Swoops in to spot PHI with laser precision!"""
@@ -80,14 +77,15 @@ def create_ad_hoc_deny_list_recognizer(deny_list: list[str] = None) -> PatternRe
     return PatternRecognizer(supported_entity="GENERIC_PII", deny_list=deny_list)
 
 def save_pdf(pdf_input) -> str:
-    """💾 Drops PDFs onto disk like hot cakes!"""
-    original_name = pdf_input.name
-    with open(original_name, "wb") as f:
-        f.write(pdf_input.read())
-    return original_name
+    """💾 Drops PDFs into a cozy temp hideout!"""
+    if pdf_input.size > 200 * 1024 * 1024:
+        raise ValueError("PDF exceeds 200MB limit")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir="/tmp") as tmp:
+        tmp.write(pdf_input.read())
+        return tmp.name
 
 # Feature Spotlight: 📄 PDF Magic Unleashed!
-# Upload a PDF, zap the PHI, and grab a shiny new file—all with a timestamp swagger! ✨
+# Zap PHI from PDFs and sling back a shiny, safe file with timestamp swagger! ✨
 
 def read_pdf(pdf_path: str) -> str:
     """📖 Slurps up PDF text like a thirsty camel!"""
@@ -111,13 +109,13 @@ model_list = [
     ("HuggingFace/obi/deid_roberta_i2b2", "https://huggingface.co/obi/deid_roberta_i2b2"),
     ("HuggingFace/StanfordAIMI/stanford-deidentifier-base", "https://huggingface.co/StanfordAIMI/stanford-deidentifier-base"),
 ]
-st_model = st.sidebar.selectbox("NER model package", [model[0] for model in model_list], index=0, help="Pick your PHI-hunting hero!")
-st.sidebar.markdown(f"[View model on HuggingFace]({next(url for model, url in model_list if model == st_model)})")
+st_model = st.sidebar.selectbox("NER model package", [m[0] for m in model_list], 0, help="Pick your PHI-hunting hero!")
+st.sidebar.markdown(f"[View model on HuggingFace]({next(url for m, url in model_list if m == st_model)})")
 st_model_package = st_model.split("/")[0]
 st_model = st_model if st_model_package.lower() != "huggingface" else "/".join(st_model.split("/")[1:])
 analyzer_params = (st_model_package, st_model)
 st.sidebar.warning("Models may take a sec to wake up!")
-st_operator = st.sidebar.selectbox("De-identification approach", ["replace", "redact", "mask"], index=0, help="Choose how to zap PHI!")
+st_operator = st.sidebar.selectbox("De-identification approach", ["replace", "redact", "mask"], 0, help="Choose how to zap PHI!")
 st_threshold = st.sidebar.slider("Acceptance threshold", 0.0, 1.0, 0.35)
 st_return_decision_process = st.sidebar.checkbox("Add analysis explanations", False)
 with st.sidebar.expander("Allowlists and denylists"):
@@ -132,8 +130,6 @@ with col1:
     if uploaded_file:
         try:
             pdf_path = save_pdf(uploaded_file)
-            if not pdf_path:
-                raise ValueError("PDF save flopped!")
             text = read_pdf(pdf_path)
             if not text:
                 raise ValueError("No text in that PDF!")
@@ -157,8 +153,6 @@ with col1:
             timestamp = get_timestamp_prefix()
             output_filename = f"{timestamp}_{uploaded_file.name}"
             pdf_output = create_pdf(anonymized_result.text, pdf_path, output_filename)
-            if not pdf_output:
-                raise ValueError("PDF creation tanked!")
             with open(output_filename, "rb") as f:
                 pdf_bytes = f.read()
                 b64 = base64.b64encode(pdf_bytes).decode()