Update app.py

app.py

@@ -17,122 +17,145 @@ dotenv.load_dotenv()
 logger = logging.getLogger("presidio-streamlit")
 
 def get_timestamp_prefix() -> str:
-    """🕒 Stamps time like a boss with Central flair!"""
+    """🕒 Stamps time with Central swagger!"""
     central = pytz.timezone("US/Central")
     return datetime.now(central).strftime("%I%M%p_%d-%m-%y").upper()
 
-def nlp_engine_and_registry(model_family: str, model_path: str) -> tuple[object, RecognizerRegistry]:
-    """🤖 Fires up NLP engines with a spark of genius!"""
+def nlp_engine_and_registry(model_family: str, model_path: str) -> tuple:
+    """🤖 Sparks NLP models with a wink!"""
     registry = RecognizerRegistry()
     if model_family.lower() == "flair":
         from flair.models import SequenceTagger
         tagger = SequenceTagger.load(model_path)
         registry.load_predefined_recognizers()
-        registry.add_recognizer_from_dict({"name": "flair_recognizer", "supported_language": "en", "supported_entities": ["PERSON", "LOCATION", "ORGANIZATION"], "model": model_path, "package": "flair"})
+        recognizer = PatternRecognizer(supported_entity="CUSTOM", supported_language="en")
+        registry.add_recognizer(recognizer)
+        logger.info(f"Flair model loaded: {model_path}")
         return tagger, registry
     elif model_family.lower() == "huggingface":
         from transformers import pipeline
         nlp = pipeline("ner", model=model_path, tokenizer=model_path)
         registry.load_predefined_recognizers()
-        registry.add_recognizer_from_dict({"name": "huggingface_recognizer", "supported_language": "en", "supported_entities": ["PERSON", "LOCATION", "ORGANIZATION", "DATE_TIME"], "model": model_path, "package": "transformers"})
+        recognizer = PatternRecognizer(supported_entity="CUSTOM", supported_language="en")
+        registry.add_recognizer(recognizer)
+        logger.info(f"HuggingFace model loaded: {model_path}")
         return nlp, registry
-    raise ValueError(f"Model family {model_family} not supported")
+    raise ValueError(f"Model family {model_family} unsupported")
 
 def analyzer_engine(model_family: str, model_path: str) -> AnalyzerEngine:
-    """🔍 Unleashes the PHI-sniffing bloodhound!"""
+    """🔍 Unleashes the PHI-hunting beast!"""
     nlp_engine, registry = nlp_engine_and_registry(model_family, model_path)
     return AnalyzerEngine(registry=registry)
 
 def get_supported_entities(model_family: str, model_path: str) -> list[str]:
-    """📋 Lists what secrets we’re hunting—PHI beware!"""
+    """📋 Spills the beans on PHI targets!"""
     return ["PERSON", "LOCATION", "ORGANIZATION", "DATE_TIME"] if model_family.lower() == "huggingface" else ["PERSON", "LOCATION", "ORGANIZATION"]
 
-# Feature Spotlight: 🕵️‍♂️ The Great PHI Hunt Begins!
-# Summon models to sniff out sensitive data in PDFs with ninja stealth! 😎
+# Feature Spotlight: 🕵️‍♂️ PHI Hunt Kicks Off!
+# Models dive into PDFs, sniffing out sensitive bits with ninja vibes! 😎
 
-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!"""
+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:
+    """🦸 Zaps PHI with eagle-eye precision!"""
     results = analyzer.analyze(text=text, entities=entities, language=language, score_threshold=score_threshold, return_decision_process=return_decision_process)
-    filtered_results = []
+    filtered = []
     for result in results:
-        text_snippet = text[result.start:result.end].lower()
-        if any(word.lower() in text_snippet for word in allow_list):
+        snippet = text[result.start:result.end].lower()
+        if any(word.lower() in snippet for word in allow_list):
             continue
-        if any(word.lower() in text_snippet for word in deny_list) or not deny_list:
-            filtered_results.append(result)
-    return filtered_results
+        if any(word.lower() in snippet for word in deny_list) or not deny_list:
+            filtered.append(result)
+    return filtered
 
-def anonymize(text: str, operator: str, analyze_results: list[RecognizerResult], mask_char: str = "*", number_of_chars: int = 15) -> dict:
-    """🕵️‍♀️ Cloaks PHI in a disguise—poof, it’s gone!"""
+def anonymize(text: str, operator: str, analyze_results: list, mask_char: str = "*", number_of_chars: int = 15) -> dict:
+    """🕵️‍♀️ Hides PHI with a magician’s flair!"""
     anonymizer = AnonymizerEngine()
-    operator_config = {"DEFAULT": OperatorConfig(operator, {})}
+    config = {"DEFAULT": OperatorConfig(operator, {})}
     if operator == "mask":
-        operator_config["DEFAULT"] = OperatorConfig(operator, {"masking_char": mask_char, "chars_to_mask": number_of_chars})
-    return anonymizer.anonymize(text=text, analyzer_results=analyze_results, operators=operator_config)
+        config["DEFAULT"] = OperatorConfig(operator, {"masking_char": mask_char, "chars_to_mask": number_of_chars})
+    return anonymizer.anonymize(text=text, analyzer_results=analyze_results, operators=config)
 
 def create_ad_hoc_deny_list_recognizer(deny_list: list[str] = None) -> PatternRecognizer:
-    """🚨 Builds a naughty list to catch sneaky PHI!"""
-    if not deny_list:
-        return None
-    return PatternRecognizer(supported_entity="GENERIC_PII", deny_list=deny_list)
+    """🚨 Sets traps for sneaky PHI rogues!"""
+    return None if not deny_list else PatternRecognizer(supported_entity="GENERIC_PII", deny_list=deny_list)
 
 def save_pdf(pdf_input) -> str:
-    """💾 Drops PDFs into a cozy temp hideout!"""
+    """💾 Stashes PDFs in a temp vault!"""
     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
+        logger.error(f"Upload rejected: {pdf_input.name} exceeds 200MB")
+        st.error("PDF exceeds 200MB limit")
+        raise ValueError("PDF too big")
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir="/tmp") as tmp:
+            tmp.write(pdf_input.read())
+            logger.info(f"Uploaded PDF to {tmp.name}, size: {pdf_input.size} bytes")
+            return tmp.name
+    except Exception as e:
+        logger.error(f"Upload failed: {str(e)}")
+        st.error(f"Upload error: {str(e)}")
+        raise
 
-# Feature Spotlight: 📄 PDF Magic Unleashed!
-# Zap PHI from PDFs and sling back a shiny, safe file with timestamp swagger! ✨
+# Feature Spotlight: 📄 PDF Wizardry Unleashed!
+# Uploads zip through, PHI vanishes, and out pops a safe PDF with timestamp pizzazz! ✨
 
 def read_pdf(pdf_path: str) -> str:
-    """📖 Slurps up PDF text like a thirsty camel!"""
-    reader = PdfReader(pdf_path)
-    return "".join(page.extract_text() or "" + "\n" for page in reader.pages)
+    """📖 Gobbles PDF text like candy!"""
+    try:
+        reader = PdfReader(pdf_path)
+        text = "".join(page.extract_text() or "" + "\n" for page in reader.pages)
+        logger.info(f"Extracted {len(text)} chars from {pdf_path}")
+        return text
+    except Exception as e:
+        logger.error(f"Read failed: {str(e)}")
+        raise
 
 def create_pdf(text: str, input_path: str, output_filename: str) -> str:
-    """🖨️ Crafts a fresh PDF with PHI-proof swagger!"""
-    reader = PdfReader(input_path)
-    writer = PdfWriter()
-    for page in reader.pages:
-        writer.add_page(page)
-    with open(output_filename, "wb") as f:
-        writer.write(f)
-    return output_filename
+    """🖨️ Spins a new PDF with PHI-proof charm!"""
+    try:
+        reader = PdfReader(input_path)
+        writer = PdfWriter()
+        for page in reader.pages:
+            writer.add_page(page)
+        with open(output_filename, "wb") as f:
+            writer.write(f)
+        logger.info(f"Created PDF: {output_filename}")
+        return output_filename
+    except Exception as e:
+        logger.error(f"Create failed: {str(e)}")
+        raise
 
-# Sidebar setup
+# Sidebar
 st.sidebar.header("PHI De-identification with Presidio")
 model_list = [
     ("flair/ner-english-large", "https://huggingface.co/flair/ner-english-large"),
     ("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", [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 = st.sidebar.selectbox("NER model", [m[0] for m in model_list], 0)
+st.sidebar.markdown(f"[View model]({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"], 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"):
-    st_allow_list = st_tags(label="Add words to allowlist", text="Enter word and press enter.")
-    st_deny_list = st_tags(label="Add words to denylist", text="Enter word and press enter.")
+st.sidebar.warning("Models may snooze briefly!")
+st_operator = st.sidebar.selectbox("De-id approach", ["replace", "redact", "mask"], 0)
+st_threshold = st.sidebar.slider("Threshold", 0.0, 1.0, 0.35)
+st_return_decision_process = st.sidebar.checkbox("Show analysis", False)
+with st.sidebar.expander("Allow/Deny lists"):
+    st_allow_list = st_tags(label="Allowlist", text="Add word, hit enter")
+    st_deny_list = st_tags(label="Denylist", text="Add word, hit enter")
 
-# Main panel
+# Main
 col1, col2 = st.columns(2)
 with col1:
     st.subheader("Input")
-    uploaded_file = st.file_uploader("Upload PDF", type=["pdf"])
+    uploaded_file = st.file_uploader("Upload PDF", type=["pdf"], help="Max 200MB")
     if uploaded_file:
         try:
+            logger.info(f"Upload: {uploaded_file.name}, size: {uploaded_file.size} bytes")
             pdf_path = save_pdf(uploaded_file)
             text = read_pdf(pdf_path)
             if not text:
-                raise ValueError("No text in that PDF!")
+                st.error("No text extracted")
+                raise ValueError("Empty PDF")
             analyzer = analyzer_engine(*analyzer_params)
             st_analyze_results = analyze(
                 analyzer=analyzer,
@@ -146,33 +169,30 @@ with col1:
             )
             phi_types = set(res.entity_type for res in st_analyze_results)
             if phi_types:
-                st.success(f"Removed PHI types: {', '.join(phi_types)}")
+                st.success(f"Zapped PHI: {', '.join(phi_types)}")
             else:
-                st.info("No PHI detected")
+                st.info("No PHI found")
             anonymized_result = anonymize(text=text, operator=st_operator, analyze_results=st_analyze_results)
             timestamp = get_timestamp_prefix()
             output_filename = f"{timestamp}_{uploaded_file.name}"
-            pdf_output = create_pdf(anonymized_result.text, pdf_path, output_filename)
+            create_pdf(anonymized_result.text, pdf_path, output_filename)
             with open(output_filename, "rb") as f:
-                pdf_bytes = f.read()
-                b64 = base64.b64encode(pdf_bytes).decode()
+                b64 = base64.b64encode(f.read()).decode()
                 st.markdown(f'<a href="data:application/pdf;base64,{b64}" download="{output_filename}">Download de-identified PDF</a>', unsafe_allow_html=True)
             with col2:
                 st.subheader("Findings")
                 if st_analyze_results:
-                    df = pd.DataFrame.from_records([r.to_dict() for r in st_analyze_results])
-                    df["text"] = [text[res.start:res.end] for res in st_analyze_results]
+                    df = pd.DataFrame([r.to_dict() for r in st_analyze_results])
+                    df["text"] = [text[r.start:r.end] for r in st_analyze_results]
                     df_subset = df[["entity_type", "text", "start", "end", "score"]].rename(
-                        {"entity_type": "Entity type", "text": "Text", "start": "Start", "end": "End", "score": "Confidence"}, axis=1
+                        {"entity_type": "Type", "text": "Text", "start": "Start", "end": "End", "score": "Confidence"}, axis=1
                     )
                     if st_return_decision_process:
-                        analysis_explanation_df = pd.DataFrame.from_records([r.analysis_explanation.to_dict() for r in st_analyze_results])
-                        df_subset = pd.concat([df_subset, analysis_explanation_df], axis=1)
+                        df_subset = pd.concat([df_subset, pd.DataFrame([r.analysis_explanation.to_dict() for r in st_analyze_results])], axis=1)
                     st.dataframe(df_subset.reset_index(drop=True), use_container_width=True)
                 else:
                     st.text("No findings")
-            if os.path.exists(pdf_path):
-                os.remove(pdf_path)
+            os.remove(pdf_path)
         except Exception as e:
-            st.error(f"Oops, something broke: {str(e)}")
-            logger.error(f"Processing error: {str(e)}")
+            st.error(f"Oops: {str(e)}")
+            logger.error(f"Error: {str(e)}")