Update app.py

app.py

@@ -8,11 +8,15 @@ 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',
@@ -26,27 +30,11 @@ DEFAULT_PRESIDIO_ENTITIES = sorted([
     'IBAN_CODE',
     'EMAIL',
 ])
-WARNING_PRESIDIO_ENTITIES = sorted([
-    'PHONE_NUMBER',
-    'US_PASSPORT',
-    'EMAIL_ADDRESS',
-    'IP_ADDRESS',
-    'US_BANK_NUMBER',
-    'IBAN_CODE',
-    'EMAIL',
-])
-ALERT_PRESIDIO_ENTITIES = sorted([
-    'CREDIT_CARD',
-    'US_SSN',
-    'US_PASSPORT',
-    'US_BANK_NUMBER',
-    'IBAN_CODE',
-])
 
 def stream_rows(dataset: str, config: str, split: str) -> Iterable[dict[str, Any]]:
     batch_size = 100
     for i in count():
-        rows_resp = requests.get(f"https://datasets-server.huggingface.co/rows?dataset={dataset}&config={config}&split={split}&offset={i * batch_size}&length={batch_size}", timeout=20).json()
+        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"]:
@@ -75,7 +63,7 @@ def presidio_report(presidio_entities: list[PresidioEntity], next_row_idx: int,
 
 
 def analyze_dataset(dataset: str, enabled_presidio_entities: list[str] = DEFAULT_PRESIDIO_ENTITIES, show_texts_without_masks: bool = False) -> pd.DataFrame:
-    info_resp = requests.get(f"https://datasets-server.huggingface.co/info?dataset={dataset}", timeout=3).json()
+    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