Update app.py

app.py

@@ -2,7 +2,7 @@ from dotenv import load_dotenv
 import os
 
 import pandas as pd
-from httpx import AsyncClient
+from httpx import Client
 from huggingface_hub import dataset_info
 from huggingface_hub.utils import logging
 from functools import lru_cache
@@ -10,28 +10,23 @@ from tqdm.contrib.concurrent import thread_map
 from huggingface_hub import HfApi
 from rich import print
 import gradio as gr
-import asyncio
 
 
-async def check_dataset_has_non_default_file(hub_id):
-    try:
-        info = await dataset_info(hub_id)
-        if files := info.siblings:
-            file_names = [f.rfilename for f in files]
-            files = [f for f in file_names if not f.startswith(".") or f == "README.md"]
-            return len(files) >= 1
-        return False
-    except Exception as e:
-        logger.error(f"Failed to get siblings for {hub_id}: {e}")
-        return False
+def check_dataset_has_non_default_file(hub_id):
+    info = dataset_info(hub_id)
+    if files := info.siblings:
+        file_names = [f.rfilename for f in files]
+        files = [f for f in file_names if not f.startswith(".") or f == "README.md"]
+        return len(files) >= 1
+    return False
 
 
-async def datasets_server_valid_rows(hub_id: str, async_client: AsyncClient):
+def datasets_server_valid_rows(hub_id: str):
     try:
-        resp = await async_client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={hub_id}")
+        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={hub_id}")
         return resp.json()["viewer"]
-    except Exception as e:
-        logger.error(f"Failed to get is-valid for {hub_id}: {e}")
+    except Exception:
+        # logger.error(f"Failed to get is-valid for {hub_id}: {e}")
         return None
 
 
@@ -48,18 +43,20 @@ BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
 
 logger = logging.get_logger(__name__)
 headers = {
-    "authorization": f"Bearer {HF_TOKEN}",
+    "authorization": f"Bearer ${HF_TOKEN}",
     "user-agent": USER_AGENT,
 }
-async_client = AsyncClient(headers=headers)
+client = Client(headers=headers)
+async_client = Client(headers=headers)
 api = HfApi(token=HF_TOKEN)
 
 
-async def get_first_config_and_split_name(hub_id: str, async_client: AsyncClient):
+def get_first_config_and_split_name(hub_id: str):
     try:
-        resp = await async_client.get(
+        resp = client.get(
             f"https://datasets-server.huggingface.co/splits?dataset={hub_id}"
         )
+
         data = resp.json()
         return data["splits"][0]["config"], data["splits"][0]["split"]
     except Exception as e:
@@ -67,29 +64,57 @@ async def get_first_config_and_split_name(hub_id: str, async_client: AsyncClient
         return None
 
 
-async def get_dataset_info(hub_id: str, config: str | None = None):
+def check_dataset_has_non_default_file(hub_id):
+    try:
+        info = dataset_info(hub_id)
+        if files := info.siblings:
+            file_names = [f.rfilename for f in files]
+            files = [f for f in file_names if not f.startswith(".") or f == "README.md"]
+            return len(files) >= 1
+        return False
+    except Exception as e:
+        logger.error(f"Failed to get siblings for {hub_id}: {e}")
+        return False
+
+
+def datasets_server_valid_rows(hub_id: str):
+    try:
+        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={hub_id}")
+        return resp.json()["viewer"]
+    except Exception:
+        # logger.error(f"Failed to get is-valid for {hub_id}: {e}")
+        return None
+
+
+def dataset_is_valid(dataset):
+    return dataset if datasets_server_valid_rows(dataset.id) else None
+
+
+def get_dataset_info(hub_id: str, config: str | None = None):
     if config is None:
-        config = await get_first_config_and_split_name(hub_id, async_client)
+        config = get_first_config_and_split_name(hub_id)
         if config is None:
             return None
         else:
             config = config[0]
-    resp = await async_client.get(
+    resp = client.get(
         f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}"
     )
     resp.raise_for_status()
     return resp.json()
 
 
-async def dataset_with_info(dataset):
+def dataset_with_info(dataset):
     try:
-        if info := await get_dataset_info(dataset.id):
+        if info := get_dataset_info(dataset.id):
             columns = info.get("dataset_info", {}).get("features", {})
             if columns is not None:
                 return {
                     "hub_id": dataset.id,
                     "column_names": list(columns.keys()),
                     "columns": columns,
+                    # "dataset": dataset,
+                    # "full_info": info,
                     "likes": dataset.likes,
                     "downloads": dataset.downloads,
                     "created_at": dataset.created_at,
@@ -100,34 +125,36 @@ async def dataset_with_info(dataset):
         return None
 
 
-async def return_dataset_with_non_default_files(dataset):
-    return dataset if await check_dataset_has_non_default_file(dataset.id) else None
+def return_dataset_with_non_default_files(dataset):
+    return dataset if check_dataset_has_non_default_file(dataset.id) else None
 
 
 @lru_cache(maxsize=100)
-async def prep_data():
-    datasets = list(api.list_datasets(limit=200, sort="createdAt", direction=-1))
+def prep_data():
+    datasets = list(api.list_datasets(limit=None, sort="createdAt", direction=-1))
     print(f"Found {len(datasets)} datasets.")
-    
-    valid_datasets = await asyncio.gather(*[return_dataset_with_non_default_files(dataset) for dataset in datasets])
-    valid_datasets = [x for x in valid_datasets if x is not None]
-    print(f"Found {len(valid_datasets)} datasets with non-default files.")
-    
-    has_server = await asyncio.gather(*[datasets_server_valid_rows(dataset.id, async_client) for dataset in valid_datasets])
-    datasets_with_server = [dataset for dataset, server_valid in zip(valid_datasets, has_server) if server_valid]
+    # datasets = thread_map(
+    #     return_dataset_with_non_default_files,
+    #     datasets,
+    # )
+    # datasets = [x for x in datasets if x is not None]
+    # print(f"Found {len(datasets)} datasets with non-default files.")
+    has_server = thread_map(
+        dataset_is_valid,
+        datasets,
+    )
+    datasets_with_server = [x for x in has_server if x is not None]
     print(f"Found {len(datasets_with_server)} datasets with server.")
-    
-    datasets_server_data = await asyncio.gather(*[dataset_with_info(dataset) for dataset in datasets_with_server])
-    datasets_server_data = [data for data in datasets_server_data if data is not None]
+    datasets_server_data = thread_map(dataset_with_info, datasets_with_server)
     print(f"Found {len(datasets_server_data)} datasets with server data.")
     print(datasets_server_data[0])
-    
     return datasets_server_data
 
 
 def filter_columns(datasets_server_data, columns=None):
     if columns is not None:
         clean = []
+        # check for presence of columns
         for dataset in datasets_server_data:
             if dataset is not None:
                 target_column = dataset.get("columns", [])
@@ -139,8 +166,17 @@ def filter_columns(datasets_server_data, columns=None):
     return datasets_server_data
 
 
-async def predict(columns_to_filter):
-    datasets_server_data = await prep_data()
+# warm up the cache
+prep_data()
+
+
+def render_model_hub_link(hub_id):
+    link = f"https://huggingface.co/datasets/{hub_id}"
+    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{hub_id}</a>'
+
+
+def predict(columns_to_filter):
+    datasets_server_data = prep_data()
     columns_to_filter = columns_to_filter.split(",")
     columns_to_filter = [x.strip() for x in columns_to_filter]
     filtered = filter_columns(
@@ -152,11 +188,6 @@ async def predict(columns_to_filter):
     return df
 
 
-def render_model_hub_link(hub_id):
-    link = f"https://huggingface.co/datasets/{hub_id}"
-    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{hub_id}</a>'
-
-
 with gr.Blocks() as demo:
     gr.Markdown("# Search Hugging Face datasets by column names (POC)")
     gr.Markdown(