missing args + set limits

app.py

@@ -1,4 +1,4 @@
-from itertools import count
+from itertools import count, islice
 from typing import Any, Iterable
 
 import gradio as gr
@@ -9,7 +9,10 @@ from gradio_huggingfacehub_search import HuggingfaceHubSearch
 
 from analyze import get_column_description, get_columns_with_strings, presidio_scan_entities
 
-def stream_rows() -> Iterable[dict[str, Any]]:
+MAX_ENTITIES = 100
+MAX_ROWS = 100
+
+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()
@@ -32,11 +35,11 @@ def analyze_dataset(dataset: str) -> pd.DataFrame:
     columns_descriptions = [
         get_column_description(column_name, features[column_name]) for column_name in scanned_columns
     ]
-    rows = stream_rows(dataset, config, split)
+    rows = islice(stream_rows(dataset, config, split), MAX_ROWS)
     presidio_entities = []
-    for presidio_entity in presidio_scan_entities(
+    for presidio_entity in islice(presidio_scan_entities(
         rows, scanned_columns=scanned_columns, columns_descriptions=columns_descriptions
-    ):
+    ), MAX_ENTITIES):
         presidio_entities.append(presidio_entity)
         yield f"Presidio scan results for {dataset}:", pd.DataFrame(presidio_entities)