update views

app.py

@@ -1,5 +1,4 @@
 import asyncio
-import json
 import re
 from typing import Dict, List
 
@@ -7,17 +6,20 @@ import gradio as gr
 import httpx
 from cashews import cache
 from huggingface_hub import ModelCard
-from ragatouille_search import create_ragatouille_interface, search_with_ragatouille
+
+from ragatouille_search import create_ragatouille_interface
 
 cache.setup("mem://")
-API_URL = "https://davanstrien-huggingface-datasets-search-v2.hf.space/similar"
+API_URL = "https://davanstrien-huggingface-datasets-search-v2.hf.space"
 HF_API_URL = "https://huggingface.co/api/datasets"
 README_URL_TEMPLATE = "https://huggingface.co/datasets/{}/raw/main/README.md"
 
 
 async def fetch_similar_datasets(dataset_id: str, limit: int = 10) -> List[Dict]:
     async with httpx.AsyncClient() as client:
-        response = await client.get(f"{API_URL}?dataset_id={dataset_id}&n={limit + 1}")
+        response = await client.get(
+            f"{API_URL}/similar?dataset_id={dataset_id}&n={limit + 1}"
+        )
         if response.status_code == 200:
             results = response.json()["results"]
             # Remove the input dataset from the results
@@ -26,9 +28,9 @@ async def fetch_similar_datasets(dataset_id: str, limit: int = 10) -> List[Dict]
 
 
 async def fetch_similar_datasets_by_text(query: str, limit: int = 10) -> List[Dict]:
-    async with httpx.AsyncClient() as client:
-        response = await client.post(
-            f"{API_URL}-by-text", params={"query": query, "n": limit + 1}
+    async with httpx.AsyncClient(timeout=30) as client:
+        response = await client.get(
+            f"{API_URL}/similar-text", params={"query": query, "n": limit + 1}
         )
         if response.status_code == 200:
             results = response.json()["results"]
@@ -150,6 +152,37 @@ async def search_similar_datasets(dataset_id: str, limit: int = 10):
     return format_results(results, dataset_cards, dataset_infos)
 
 
+async def search_viewer(query: str, limit: int = 10):
+    async with httpx.AsyncClient(timeout=30) as client:
+        response = await client.get(
+            f"{API_URL}/search-viewer", params={"query": query, "n": limit}
+        )
+        if response.status_code == 200:
+            results = response.json()["results"]
+            return format_viewer_results(results)
+        return "No results found."
+
+
+def format_viewer_results(results: List[Dict]) -> str:
+    html = "<div style='height: 600px; overflow-y: auto;'>"
+    for result in results:
+        dataset_id = result["dataset_id"]
+        html += f"""
+        <div style='margin-bottom: 20px; border: 1px solid #ddd; padding: 10px;'>
+            <h3>{dataset_id}</h3>
+            <p><strong>Similarity Score:</strong> {result['similarity']:.4f}</p>
+            <iframe
+                src="https://huggingface.co/datasets/{dataset_id}/embed/viewer/default/train"
+                frameborder="0"
+                width="100%"
+                height="560px"
+            ></iframe>
+        </div>
+        """
+    html += "</div>"
+    return html
+
+
 with gr.Blocks() as demo:
     gr.Markdown("## &#129303; Dataset Search and Similarity")
 
@@ -219,4 +252,35 @@ with gr.Blocks() as demo:
         with gr.TabItem("RAGatouille Search"):
             ragatouille_interface = create_ragatouille_interface()
 
+        with gr.TabItem("Search Viewer"):
+            gr.Markdown("## &#128269; Search Viewer")
+            with gr.Row():
+                gr.Markdown(
+                    "This tab allows you to search for datasets using the Search Viewer endpoint. "
+                    "Enter a query to find relevant datasets and preview them."
+                )
+
+            with gr.Row():
+                viewer_query = gr.Textbox(
+                    label="Search Query", placeholder="Enter your search query here"
+                )
+
+            with gr.Row():
+                viewer_search_btn = gr.Button("Search")
+                viewer_max_results = gr.Slider(
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=10,
+                    label="Maximum number of results",
+                )
+
+            viewer_results = gr.HTML()
+
+            viewer_search_btn.click(
+                lambda query, limit: asyncio.run(search_viewer(query, limit)),
+                inputs=[viewer_query, viewer_max_results],
+                outputs=viewer_results,
+            )
+
 demo.launch()