davidberenstein1957's picture
feat: small ux updates
e30a182
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()