Update app.py

app.py

@@ -11,6 +11,19 @@ from analyze import analyzer, get_column_description, get_columns_with_strings,
 
 MAX_ROWS = 100
 T = TypeVar("T")
+DEFAULT_PRESIDIO_ENTITIES = sorted([
+    'PERSON',
+    'CREDIT_CARD',
+    'US_SSN',
+    'US_DRIVER_LICENSE',
+    'PHONE_NUMBER',
+    'US_PASSPORT',
+    'EMAIL_ADDRESS',
+    'IP_ADDRESS',
+    'US_BANK_NUMBER',
+    'IBAN_CODE',
+    'EMAIL',
+])
 
 def stream_rows(dataset: str, config: str, split: str) -> Iterable[dict[str, Any]]:
     batch_size = 100
@@ -52,7 +65,7 @@ def analyze_dataset(dataset: str, enabled_presidio_entities: str) -> pd.DataFram
     for presidio_entity in presidio_scan_entities(
         rows, scanned_columns=scanned_columns, columns_descriptions=columns_descriptions
     ):
-        if presidio_entity.type in enabled_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)
@@ -68,8 +81,8 @@ with gr.Blocks() as demo:
         ),
         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"],
+            choices=sorted(analyzer.get_supported_entities()),
+            value=DEFAULT_PRESIDIO_ENTITIES,
             interative=True,
         ),
     ]
@@ -80,7 +93,14 @@ with gr.Blocks() as demo:
     ]
     button.click(analyze_dataset, inputs, outputs)
     gr.Examples(
-        [["microsoft/orca-math-word-problems-200k"], ["tatsu-lab/alpaca"], ["Anthropic/hh-rlhf"], ["OpenAssistant/oasst1"], ["sidhq/email-thread-summary"], ["lhoestq/fake_name_and_ssn"]],
+        [
+            ["microsoft/orca-math-word-problems-200k", DEFAULT_PRESIDIO_ENTITIES],
+            ["tatsu-lab/alpaca", DEFAULT_PRESIDIO_ENTITIES],
+            ["Anthropic/hh-rlhf", DEFAULT_PRESIDIO_ENTITIES],
+            ["OpenAssistant/oasst1", DEFAULT_PRESIDIO_ENTITIES],
+            ["sidhq/email-thread-summary", DEFAULT_PRESIDIO_ENTITIES],
+            ["lhoestq/fake_name_and_ssn", DEFAULT_PRESIDIO_ENTITIES]
+        ],
         inputs,
         outputs,
         fn=analyze_dataset,