Update app.py

app.py

@@ -1,34 +1,57 @@
-"""Streamlit app for Presidio."""
 
-import json
-from json import JSONEncoder
-from annotated_text import annotated_text
-import pandas as pd
-import streamlit as st
-from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
-from presidio_anonymizer import AnonymizerEngine
-
-from flair_recognizer import FlairRecognizer
+"""Streamlit app for Presidio + Privy-trained PII models."""
 
 import spacy
-spacy.cli.download("en_core_web_lg")
+from spacy_recognizer import CustomSpacyRecognizer
+from presidio_analyzer.nlp_engine import NlpEngineProvider
+from presidio_anonymizer import AnonymizerEngine
+from presidio_analyzer import AnalyzerEngine, RecognizerRegistry
+import pandas as pd
+from annotated_text import annotated_text
+from json import JSONEncoder
+import json
+import warnings
+import streamlit as st
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+warnings.filterwarnings('ignore')
+# from flair_recognizer import FlairRecognizer
 
 # Helper methods
 @st.cache(allow_output_mutation=True)
 def analyzer_engine():
     """Return AnalyzerEngine."""
 
-    flair_recognizer = FlairRecognizer()
-    
+    spacy_recognizer = CustomSpacyRecognizer()
+
+    configuration = {
+        "nlp_engine_name": "spacy",
+        "models": [
+            {"lang_code": "en", "model_name": "en_spacy_pii_distilbert"}],
+    }
+
+    # Create NLP engine based on configuration
+    provider = NlpEngineProvider(nlp_configuration=configuration)
+    nlp_engine = provider.create_engine()
+
     registry = RecognizerRegistry()
-    registry.add_recognizer(flair_recognizer)
-    registry.load_predefined_recognizers()
+    # add rule-based recognizers
+    registry.load_predefined_recognizers(nlp_engine=nlp_engine)
+    registry.add_recognizer(spacy_recognizer)
+    # remove the nlp engine we passed, to use custom label mappings
     registry.remove_recognizer("SpacyRecognizer")
-    
-    analyzer = AnalyzerEngine(registry=registry)
+
+    analyzer = AnalyzerEngine(nlp_engine=nlp_engine,
+                              registry=registry, supported_languages=["en"])
+
+    # uncomment for flair-based NLP recognizer
+    # flair_recognizer = FlairRecognizer()
+    # registry.load_predefined_recognizers()
+    # registry.add_recognizer(flair_recognizer)
+    # analyzer = AnalyzerEngine(registry=registry, supported_languages=["en"])
     return analyzer
 
-    
+
 @st.cache(allow_output_mutation=True)
 def anonymizer_engine():
     """Return AnonymizerEngine."""
@@ -49,10 +72,12 @@ def analyze(**kwargs):
 
 def anonymize(text, analyze_results):
     """Anonymize identified input using Presidio Abonymizer."""
-
+    if not text:
+        return
     res = anonymizer_engine().anonymize(text, analyze_results)
     return res.text
 
+
 def annotate(text, st_analyze_results, st_entities):
     tokens = []
     # sort by start index
@@ -72,12 +97,14 @@ def annotate(text, st_analyze_results, st_entities):
             tokens.append(text[res.end:])
     return tokens
 
-st.set_page_config(page_title="Presidio demo (English)", layout="wide")
+
+st.set_page_config(page_title="Privy + Presidio demo (English)", layout="wide")
 
 # Side bar
 st.sidebar.markdown(
-"""
-Detect and anonymize PII in text using an [NLP model](https://huggingface.co/beki/en_spacy_pii_distilbert) trained on protocol trace data generated by [privy](https://github.com/pixie-io/pixie/tree/main/src/datagen/pii/privy) and rule-based classifiers from [presidio](https://aka.ms/presidio). 
+    """
+Detect and anonymize PII in text using an [NLP model](https://huggingface.co/beki/en_spacy_pii_distilbert) trained on protocol traces (JSON, SQL, XML etc.) generated by 
+[Privy](https://github.com/pixie-io/pixie/tree/main/src/datagen/pii/privy) and rule-based classifiers from [Presidio](https://aka.ms/presidio).
 """
 )
 
@@ -91,7 +118,8 @@ st_threshold = st.sidebar.slider(
     label="Acceptance threshold", min_value=0.0, max_value=1.0, value=0.35
 )
 
-st_return_decision_process = st.sidebar.checkbox("Add analysis explanations in json")
+st_return_decision_process = st.sidebar.checkbox(
+    "Add analysis explanations in json")
 
 st.sidebar.info(
     "Privy is an open source framework for synthetic data generation in protocol trace formats (json, sql, html etc). Presidio is an open source framework for PII detection and anonymization. "
@@ -100,42 +128,49 @@ st.sidebar.info(
 
 
 # Main panel
-analyzer_load_state = st.info("Starting Presidio analyzer and loading Privy-trained model...")
+analyzer_load_state = st.info(
+    "Starting Presidio analyzer and loading Privy-trained PII model...")
 engine = analyzer_engine()
 analyzer_load_state.empty()
 
 
 st_text = st.text_area(
     label="Type in some text",
-    value=
-    "SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"
+    value="SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"
     "\n\n"
     "{user: Willie Porter, ip: 192.168.2.80, email: willie@gmail.com}",
     height=200,
 )
 
+button = st.button("Detect PII")
+
+if 'first_load' not in st.session_state:
+    st.session_state['first_load'] = True
+
 # After
 st.subheader("Analyzed")
 with st.spinner("Analyzing..."):
-    st_analyze_results = analyze(
-        text=st_text,
-        entities=st_entities,
-        language="en",
-        score_threshold=st_threshold,
-        return_decision_process=st_return_decision_process,
-    )
-    annotated_tokens = annotate(st_text, st_analyze_results, st_entities)
-    # annotated_tokens
-    annotated_text(*annotated_tokens)
-
+    if button or st.session_state.first_load:
+        st_analyze_results = analyze(
+            text=st_text,
+            entities=st_entities,
+            language="en",
+            score_threshold=st_threshold,
+            return_decision_process=st_return_decision_process,
+        )
+        annotated_tokens = annotate(st_text, st_analyze_results, st_entities)
+        # annotated_tokens
+        annotated_text(*annotated_tokens)
 # vertical space
 st.text("")
-    
+
 st.subheader("Anonymized")
 
 with st.spinner("Anonymizing..."):
-    st_anonymize_results = anonymize(st_text, st_analyze_results)
-    st_anonymize_results
+    if button or st.session_state.first_load:
+        st_anonymize_results = anonymize(st_text, st_analyze_results)
+        st_anonymize_results
+
 
 # table result
 st.subheader("Detailed Findings")
@@ -155,11 +190,14 @@ if st_analyze_results:
     )
 
     st.dataframe(df, width=1000)
-# table result
 else:
     st.text("No findings")
 
+st.session_state['first_load'] = True
+
 # json result
+
+
 class ToDictListEncoder(JSONEncoder):
     """Encode dict to json."""