fix details

README.md

@@ -1,6 +1,6 @@
 ---
 title: Datasets Similarity Tool
-emoji: 🐨
+emoji: 🕵️‍♀️
 colorFrom: blue
 colorTo: purple
 sdk: gradio

app.py

@@ -1,17 +1,19 @@
-from dotenv import load_dotenv
+# Inspired by https://huggingface.co/spaces/davanstrien/dataset_column_search
+
 import os
+from functools import lru_cache
+from urllib.parse import quote
 
+import faiss
+import gradio as gr
+import numpy as np
 import pandas as pd
+from dotenv import load_dotenv
 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 huggingface_hub.utils import logging
 from sentence_transformers import SentenceTransformer
-import faiss
-import numpy as np
-from urllib.parse import quote
+from tqdm.contrib.concurrent import thread_map
 
 load_dotenv()
 
@@ -53,10 +55,7 @@ def dataset_is_valid(dataset):
 
 def get_first_config_and_split_name(hub_id: str):
     try:
-        resp = client.get(
-            f"https://datasets-server.huggingface.co/splits?dataset={hub_id}"
-        )
-
+        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:
@@ -71,9 +70,7 @@ def get_dataset_info(hub_id: str, config: str | None = None):
             return None
         else:
             config = config[0]
-    resp = client.get(
-        f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}"
-    )
+    resp = client.get(f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}")
     resp.raise_for_status()
     return resp.json()
 
@@ -91,35 +88,34 @@ def dataset_with_info(dataset):
                     "downloads": dataset.downloads,
                     "created_at": dataset.created_at,
                     "tags": dataset.tags,
+                    "text": f"{dataset.id}-{','.join(list(columns.keys()))}",
                 }
     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.")
+    print(f"Found {len(datasets_with_server)} valid datasets.")
     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])
+    print(f"Found {len(dataset_infos)} datasets with info.")
     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")
+encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 vectors = encoder.encode(text)
 vector_dimension = vectors.shape[1]
 print("Start indexing")
@@ -128,24 +124,25 @@ 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):
+def search(dataset_name, k):
     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"}])
+        return pd.DataFrame([{"error": "❌ 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]})
+    distances, ann = index.search(_vector, k=k)
+    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)
@@ -153,14 +150,14 @@ def search(dataset_name):
     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"
-    )
+    gr.Markdown("This space shows similar datasets based on a name and columns. It uses https://github.com/facebookresearch/faiss for vector indexing.")
+    dataset_name = gr.Textbox("asoria/bolivian-population", label="Dataset Name")
+    k = gr.Slider(5, 200, 20, step=5, interactive=True, label="K Nearest Neighbors")
     btn = gr.Button("Show similar datasets")
     df = gr.DataFrame(datatype="markdown")
-    btn.click(search, dataset_name, df)
+    btn.click(search, inputs=[dataset_name, k], outputs=df)
 
 demo.launch()