| 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() | |