Update app.py

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 = [