feat: add basic app for query composer

.gitignore

@@ -0,0 +1 @@
+.env

README.md

@@ -1,12 +1,12 @@
 ---
 title: Text To Sql Hub Datasets
-emoji: 🏆
+emoji: 🐥 🦙 🤗
 colorFrom: blue
 colorTo: blue
 sdk: gradio
 sdk_version: 4.44.0
 app_file: app.py
-pinned: false
+pinned: true
 license: apache-2.0
 ---
 

app.py

@@ -0,0 +1,116 @@
+import json
+import os
+import urllib.parse
+
+import gradio as gr
+import requests
+from gradio_huggingfacehub_search import HuggingfaceHubSearch
+from huggingface_hub import InferenceClient
+
+example = HuggingfaceHubSearch().example_value()
+
+client = InferenceClient(
+    "meta-llama/Meta-Llama-3.1-70B-Instruct",
+    token=os.environ["HF_TOKEN"],
+)
+
+
+def get_iframe(hub_repo_id, sql_query=None):
+    if sql_query:
+        sql_query = urllib.parse.quote(sql_query)
+        url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer?sql_console=true&sql={sql_query}"
+    else:
+        url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer"
+    iframe = f"""
+    <iframe
+  src="{url}"
+  frameborder="0"
+  width="100%"
+  height="560px"
+></iframe>
+"""
+    return iframe
+
+
+def get_column_info(hub_repo_id):
+    url: str = f"https://datasets-server.huggingface.co/info?dataset={hub_repo_id}"
+    response = requests.get(url)
+    data = response.json()
+    data = data.get("dataset_info")
+    key = list(data.keys())[0]
+    description = data.get(key).get("description")
+    features = json.dumps(data.get(key).get("features"))
+    return description, features
+
+
+def query_dataset(hub_repo_id, description, features, query):
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant that returns a DuckDB SQL query based on the user's query and dataset features. Only return the SQL query, no other text.",
+        },
+        {
+            "role": "user",
+            "content": f"""# Description
+{description}
+
+# Features
+{features}
+
+# Query
+{query}
+""",
+        },
+    ]
+    response = client.chat_completion(
+        messages=messages,
+        max_tokens=1000,
+        stream=False,
+    )
+    query = response.choices[0].message.content
+    return query, get_iframe(hub_repo_id, query)
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            search_in = HuggingfaceHubSearch(
+                label="Search Huggingface Hub",
+                placeholder="Search for models on Huggingface",
+                search_type="dataset",
+            )
+
+            btn = gr.Button("Show Dataset")
+    with gr.Row():
+        search_out = gr.HTML(label="Search Results")
+    with gr.Row():
+        description = gr.Textbox(
+            label="Description", placeholder="Description from dataset or project page"
+        )
+        features = gr.Code(label="Features", language="json")
+    with gr.Row():
+        query = gr.Textbox(label="Query")
+    with gr.Row():
+        sql_out = gr.Code(label="SQL Query")
+    with gr.Row():
+        btn2 = gr.Button("Query Dataset")
+
+    gr.on(
+        [btn.click, search_in.submit],
+        fn=get_iframe,
+        inputs=[search_in],
+        outputs=[search_out],
+    ).then(
+        fn=get_column_info,
+        inputs=[search_in],
+        outputs=[description, features],
+    )
+
+    btn2.click(
+        fn=query_dataset,
+        inputs=[search_in, description, features, query],
+        outputs=[sql_out],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

demo.py

@@ -0,0 +1,13 @@
+import requests
+
+
+def get_column_info(hub_repo_id):
+    url = f"https://datasets-server.huggingface.co/info?dataset={hub_repo_id}"
+    response = requests.get(url)
+    data = response.json()
+    description = data.get("dataset_info").get(hub_repo_id).get("description")
+    features = data.get("dataset_info").get(hub_repo_id).get("features")
+    return description, features
+
+
+get_column_info("davidberenstein1957/test")

requirements.txt

@@ -0,0 +1,2 @@
+gradio-huggingfacehub-search
+gradio