File size: 3,075 Bytes
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 |
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()
|