Spaces:
Sleeping
Sleeping
| from collections import Counter | |
| from itertools import count, groupby, islice | |
| from operator import itemgetter | |
| from typing import Any, Iterable, TypeVar | |
| import gradio as gr | |
| import requests | |
| import pandas as pd | |
| from datasets import Features | |
| from gradio_huggingfacehub_search import HuggingfaceHubSearch | |
| from requests.adapters import HTTPAdapter, Retry | |
| from analyze import PresidioEntity, analyzer, get_column_description, get_columns_with_strings, mask, presidio_scan_entities | |
| MAX_ROWS = 100 | |
| T = TypeVar("T") | |
| session = requests.Session() | |
| retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504]) | |
| session.mount('http://', HTTPAdapter(max_retries=retries)) | |
| 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 | |
| for i in count(): | |
| rows_resp = session.get(f"https://datasets-server.huggingface.co/rows?dataset={dataset}&config={config}&split={split}&offset={i * batch_size}&length={batch_size}", timeout=10).json() | |
| if "error" in rows_resp: | |
| raise RuntimeError(rows_resp["error"]) | |
| if not rows_resp["rows"]: | |
| break | |
| for row_item in rows_resp["rows"]: | |
| yield row_item["row"] | |
| class track_iter: | |
| def __init__(self, it: Iterable[T]): | |
| self.it = it | |
| self.next_idx = 0 | |
| def __iter__(self) -> T: | |
| for item in self.it: | |
| self.next_idx += 1 | |
| yield item | |
| def presidio_report(presidio_entities: list[PresidioEntity], next_row_idx: int, num_rows: int) -> dict[str, float]: | |
| title = f"Scan finished: {len(presidio_entities)} entities found" if num_rows == next_row_idx else "Scan in progress..." | |
| counter = Counter([title] * next_row_idx) | |
| for row_idx, presidio_entities_per_row in groupby(presidio_entities, itemgetter("row_idx")): | |
| counter.update(set("% of rows with " + presidio_entity["type"] for presidio_entity in presidio_entities_per_row)) | |
| return dict((presidio_entity_type, presidio_entity_type_row_count / num_rows) for presidio_entity_type, presidio_entity_type_row_count in counter.most_common()) | |
| def analyze_dataset(dataset: str, enabled_presidio_entities: list[str] = DEFAULT_PRESIDIO_ENTITIES, show_texts_without_masks: bool = False) -> pd.DataFrame: | |
| info_resp = session.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json() | |
| if "error" in info_resp: | |
| yield "❌ " + info_resp["error"], pd.DataFrame() | |
| return | |
| config = "default" if "default" in info_resp["dataset_info"] else next(iter(info_resp["dataset_info"])) | |
| features = Features.from_dict(info_resp["dataset_info"][config]["features"]) | |
| split = "train" if "train" in info_resp["dataset_info"][config]["splits"] else next(iter(info_resp["dataset_info"][config]["splits"])) | |
| num_rows = min(info_resp["dataset_info"][config]["splits"][split]["num_examples"], MAX_ROWS) | |
| scanned_columns = get_columns_with_strings(features) | |
| columns_descriptions = [ | |
| get_column_description(column_name, features[column_name]) for column_name in scanned_columns | |
| ] | |
| rows = track_iter(islice(stream_rows(dataset, config, split), MAX_ROWS)) | |
| presidio_entities = [] | |
| for presidio_entity in presidio_scan_entities( | |
| rows, scanned_columns=scanned_columns, columns_descriptions=columns_descriptions | |
| ): | |
| if not show_texts_without_masks: | |
| presidio_entity["text"] = mask(presidio_entity["text"]) | |
| if presidio_entity["type"] in enabled_presidio_entities: | |
| presidio_entities.append(presidio_entity) | |
| yield presidio_report(presidio_entities, next_row_idx=rows.next_idx, num_rows=num_rows), pd.DataFrame(presidio_entities) | |
| yield presidio_report(presidio_entities, next_row_idx=rows.next_idx, num_rows=num_rows), pd.DataFrame(presidio_entities) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Scan datasets using Presidio") | |
| gr.Markdown("The space takes an HF dataset name as an input, and returns the list of entities detected by Presidio in the first samples.") | |
| inputs = [ | |
| HuggingfaceHubSearch( | |
| label="Hub Dataset ID", | |
| placeholder="Search for dataset id on Huggingface", | |
| search_type="dataset", | |
| ), | |
| gr.CheckboxGroup( | |
| label="Presidio entities", | |
| choices=sorted(analyzer.get_supported_entities()), | |
| value=DEFAULT_PRESIDIO_ENTITIES, | |
| interactive=True, | |
| ), | |
| gr.Checkbox(label="Show texts without masks", value=False), | |
| ] | |
| button = gr.Button("Run Presidio Scan") | |
| outputs = [ | |
| gr.Label(show_label=False), | |
| gr.DataFrame(), | |
| ] | |
| 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"] | |
| ], | |
| inputs, | |
| outputs, | |
| fn=analyze_dataset, | |
| run_on_click=True, | |
| cache_examples=False, | |
| ) | |
| demo.launch() | |