Add application file

README.md

@@ -1,12 +1,13 @@
 ---
 title: Datasets Similarity Tool
-emoji: 📉
-colorFrom: pink
-colorTo: pink
+emoji: 🐨
+colorFrom: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 4.19.2
 app_file: app.py
 pinned: false
+startup_duration_timeout: 1h
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,166 @@
+from dotenv import load_dotenv
+import os
+
+import pandas as pd
+from httpx import Client
+from huggingface_hub.utils import logging
+from functools import lru_cache
+from tqdm.contrib.concurrent import thread_map
+from huggingface_hub import HfApi
+import gradio as gr
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+from urllib.parse import quote
+
+load_dotenv()
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+assert HF_TOKEN is not None, "You need to set HF_TOKEN in your environment variables"
+
+BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
+
+logger = logging.get_logger(__name__)
+headers = {
+    "authorization": f"Bearer ${HF_TOKEN}",
+}
+client = Client(headers=headers)
+api = HfApi(token=HF_TOKEN)
+
+
+def get_first_config_name(dataset: str):
+    try:
+        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/splits?dataset={dataset}")
+        data = resp.json()
+        return data["splits"][0]["config"][0]
+    except Exception as e:
+        logger.error(f"Failed to get splits for {dataset}: {e}")
+        return None
+
+
+def datasets_server_valid_rows(dataset: str):
+    try:
+        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={dataset}")
+        return resp.json()["viewer"]
+    except Exception as e:
+        logger.error(f"Failed to get is-valid for {dataset}: {e}")
+        return None
+
+
+def dataset_is_valid(dataset):
+    return dataset if datasets_server_valid_rows(dataset.id) else None
+
+
+def get_first_config_and_split_name(hub_id: str):
+    try:
+        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:
+        logger.error(f"Failed to get splits for {hub_id}: {e}")
+        return None
+
+
+def get_dataset_info(hub_id: str, config: str | None = None):
+    if config is None:
+        config = get_first_config_and_split_name(hub_id)
+        if config is None:
+            return None
+        else:
+            config = config[0]
+    resp = client.get(
+        f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}"
+    )
+    resp.raise_for_status()
+    return resp.json()
+
+
+def dataset_with_info(dataset):
+    try:
+        if info := get_dataset_info(dataset.id):
+            columns = info.get("dataset_info", {}).get("features", {})
+            if columns is not None:
+                return {
+                    "dataset": dataset.id,
+                    "column_names": ','.join(list(columns.keys())),
+                    "text": f"{dataset.id}-{','.join(list(columns.keys()))}",
+                    "likes": dataset.likes,
+                    "downloads": dataset.downloads,
+                    "created_at": dataset.created_at,
+                    "tags": dataset.tags,
+                }
+    except Exception as e:
+        logger.error(f"Failed to get info for {dataset.id}: {e}")
+        return None
+
+
+
+@lru_cache(maxsize=100)
+def prep_data():
+    datasets = list(api.list_datasets(limit=None, sort="createdAt", direction=-1))
+    print(f"Found {len(datasets)} datasets in the hub.")
+    logger.info(f"Found {len(datasets)} datasets.")
+    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.")
+    dataset_infos = thread_map(dataset_with_info, datasets_with_server)
+    dataset_infos = [x for x in dataset_infos if x is not None]
+    print(f"Found {len(dataset_infos)} datasets with server data.")
+    print(dataset_infos[0])
+    return dataset_infos
+
+all_datasets = prep_data()
+all_datasets_df = pd.DataFrame.from_dict(all_datasets)
+print(all_datasets_df.head())
+text = all_datasets_df['text']
+encoder = SentenceTransformer("paraphrase-mpnet-base-v2")
+vectors = encoder.encode(text)
+vector_dimension = vectors.shape[1]
+print("Start indexing")
+index = faiss.IndexFlatL2(vector_dimension)
+faiss.normalize_L2(vectors)
+index.add(vectors)
+print("Indexing done")
+
+def render_model_hub_link(hub_id):
+    link = f"https://huggingface.co/datasets/{quote(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 search(dataset_name):
+    print(f"start search for {dataset_name}")
+    try:
+        dataset_row = all_datasets_df[all_datasets_df.dataset == dataset_name].iloc[0]
+        print(dataset_row)
+    except IndexError:
+        return pd.DataFrame([{"error": f"❌ Dataset does not exist or is not supported"}])
+    text = dataset_row["text"]
+    search_vector = encoder.encode(text)
+    _vector = np.array([search_vector])
+    faiss.normalize_L2(_vector)
+    distances, ann = index.search(_vector, k=20)
+    results = pd.DataFrame({'distances': distances[0], 'ann': ann[0]})
+    print("results for distances and ann")
+    print(results)
+    merge = pd.merge(results, all_datasets_df, left_on="ann", right_index=True)
+    print("resultst for merged df (distances,ann, dataset info)")
+    merge["dataset"] = merge["dataset"].apply(render_model_hub_link)
+    return merge
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Search similar Datasets on Hugging Face")
+    gr.Markdown("This space shows similar dataset based on column name and types")
+    dataset_name = gr.Textbox(
+        "asoria/bolivian-population", label="Dataset Name"
+    )
+    btn = gr.Button("Show similar datasets")
+    df = gr.DataFrame(datatype="markdown")
+    btn.click(search, dataset_name, df)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+Pyarrow
+gradio==4.18.0
+httpx
+huggingface_hub
+pandas
+python-dotenv
+datasets
+sentence-transformers
+faiss-cpu