File size: 3,604 Bytes
f6aec2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75d3b30
f6aec2d
 
 
 
 
 
 
 
07850ea
 
 
 
 
 
 
 
f6aec2d
 
07850ea
f6aec2d
 
 
 
 
 
 
07850ea
f6aec2d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0cc19e7
 
 
 
 
 
f6aec2d
 
 
 
 
 
 
 
 
 
 
 
07850ea
f6aec2d
07850ea
f6aec2d
 
 
 
 
 
 
 
 
 
 
 
 
07850ea
f6aec2d
 
 
 
07850ea
 
f6aec2d
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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="800px"
></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)
    try:
        data = response.json()
        data = data.get("dataset_info")
        key = list(data.keys())[0]
        features: str = json.dumps(data.get(key).get("features"))
    except Exception as e:
        gr.Error(f"Error getting column info: {e}")
    return features


def query_dataset(hub_repo_id, 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"""table train
# 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:
    gr.Markdown("""# πŸ₯ πŸ¦™ πŸ€— Text To Sql Hub Datasets πŸ₯ πŸ¦™ πŸ€—

                This is a basic text to SQL tool that allows you to query datasets on Huggingface Hub.
                It is built with [DuckDB](https://duckdb.org/), [Huggingface's Inference API](https://huggingface.co./docs/api-inference/index), and [LLama 3.1 70B](https://huggingface.co./meta-llama/Meta-Llama-3.1-70B-Instruct).
                Also, it uses the [dataset-server API](https://redocly.github.io/redoc/?url=https://datasets-server.huggingface.co/openapi.json#operation/isValidDataset).
                """)
    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():
        features = gr.Code(label="Features", language="json", visible=False)
    with gr.Row():
        query = gr.Textbox(label="Query", placeholder="Enter a query to generate SQL")
    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=[features],
    )

    btn2.click(
        fn=query_dataset,
        inputs=[search_in, features, query],
        outputs=[sql_out, search_out],
    )

if __name__ == "__main__":
    demo.launch()