Update app.py

analyze.py

@@ -11,7 +11,8 @@ Row = dict[str, Any]
 T = TypeVar("T")
 BATCH_SIZE = 1
 MAX_TEXT_LENGTH = 500
-batch_analyzer: Optional[BatchAnalyzerEngine] = None
+analyzer = AnalyzerEngine()
+batch_analyzer = BatchAnalyzerEngine(analyzer)
 
 
 class PresidioEntity(TypedDict):
@@ -121,16 +122,13 @@ def analyze(
 def presidio_scan_entities(
     rows: Iterable[Row], scanned_columns: list[str], columns_descriptions: list[str]
 ) -> Iterable[PresidioEntity]:
-    global batch_analyzer
     cache: dict[str, list[RecognizerResult]] = {}
-    if batch_analyzer is None:
-        batch_analyser = BatchAnalyzerEngine(AnalyzerEngine())
     rows_with_scanned_columns_only = (
         {column_name: get_strings(row[column_name])[:MAX_TEXT_LENGTH] for column_name in scanned_columns} for row in rows
     )
     for indices, batch in batched(rows_with_scanned_columns_only, BATCH_SIZE, with_indices=True):
         yield from analyze(
-            batch_analyzer=batch_analyser,
+            batch_analyzer=batch_analyzer,
             batch=batch,
             indices=indices,
             scanned_columns=scanned_columns,

app.py

@@ -7,7 +7,7 @@ import pandas as pd
 from datasets import Features
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 
-from analyze import get_column_description, get_columns_with_strings, presidio_scan_entities
+from analyze import analyzer, get_column_description, get_columns_with_strings, presidio_scan_entities
 
 MAX_ROWS = 100
 T = TypeVar("T")
@@ -34,7 +34,7 @@ class track_iter:
             self.next_idx += 1
             yield item
 
-def analyze_dataset(dataset: str) -> pd.DataFrame:
+def analyze_dataset(dataset: str, enabled_presidio_entities: str) -> pd.DataFrame:
     info_resp = requests.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json()
     if "error" in info_resp:
         yield "❌ " + info_resp["error"], pd.DataFrame()
@@ -52,8 +52,9 @@ def analyze_dataset(dataset: str) -> pd.DataFrame:
     for presidio_entity in presidio_scan_entities(
         rows, scanned_columns=scanned_columns, columns_descriptions=columns_descriptions
     ):
-        presidio_entities.append(presidio_entity)
-        yield f"⚙️ Scanning {dataset} [{rows.next_idx}/{num_rows} rows]:", pd.DataFrame(presidio_entities)
+        if presidio_entity.type in enabled_presidio_entities:
+            presidio_entities.append(presidio_entity)
+            yield f"⚙️ Scanning {dataset} [{rows.next_idx}/{num_rows} rows]:", pd.DataFrame(presidio_entities)
     yield f"✅ Scanning {dataset} [{rows.next_idx}/{num_rows} rows]:", pd.DataFrame(presidio_entities)
 
 with gr.Blocks() as demo:
@@ -65,6 +66,12 @@ with gr.Blocks() as demo:
             placeholder="Search for dataset id on Huggingface",
             search_type="dataset",
         ),
+        gr.CheckBoxGroup(
+            label="Presidio entities",
+            choices=analyzer.get_supported_entities(),
+            value=["PERSON", "CREDIT_CARD", "US_SSN", "PHONE_NUMBER", "EMAIL_ADDRESS", "IP_ADDRESS", "US_BANK_NUMBER", "EMAIL", "IBAN_CODE"],
+            interative=True,
+        ),
     ]
     button = gr.Button("Run Presidio Scan")
     outputs = [