File size: 4,122 Bytes
f6aec2d 8f8b860 f6aec2d 75d3b30 f6aec2d 07850ea f6aec2d 07850ea f6aec2d 2377601 f6aec2d 07850ea f6aec2d bde0dbe 0cc19e7 f6aec2d b92493b f6aec2d e30a182 f6aec2d e30a182 f6aec2d 07850ea f6aec2d 07850ea f6aec2d e30a182 f6aec2d 07850ea 3c508fe 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 121 122 123 124 125 126 127 128 129 130 |
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 not hub_repo_id:
raise ValueError("Hub repo id is required")
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 SQL query expert assistant that returns a DuckDB SQL queries based on the user's natural language query and dataset features. You might need to use DuckDB functions for lists and aggregations, given the 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",
sumbit_on_select=True,
)
query = gr.Textbox(
label="Natural Language Query",
placeholder="Enter a natural language query to generate SQL",
)
sql_out = gr.Code(
label="SQL Query",
interactive=True,
language="sql",
lines=1,
visible=False,
)
with gr.Row():
with gr.Column():
btn = gr.Button("Show Dataset")
with gr.Column():
btn2 = gr.Button("Query Dataset")
with gr.Row():
search_out = gr.HTML(label="Search Results")
with gr.Row():
features = gr.Code(label="Features", language="json", visible=False)
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],
)
gr.on(
[btn2.click, query.submit],
fn=query_dataset,
inputs=[search_in, features, query],
outputs=[sql_out, search_out],
)
if __name__ == "__main__":
demo.launch()
|