asoria/datasets-text2sql

1import os2import duckdb3import gradio as gr4from httpx import Client5from huggingface_hub import HfApi6import pandas as pd7from gradio_huggingfacehub_search import HuggingfaceHubSearch8import spaces9from llama_cpp import Llama10 11 12BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"13headers = {14	"Accept" : "application/json",15	"Content-Type": "application/json" 16}17client = Client(headers=headers)18api = HfApi()19llama = Llama(20        model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",21        n_ctx=2048,22        n_gpu_layers=5023    )24 25@spaces.GPU26def generate_sql(prompt):27    # pred = pipe(prompt, max_length=1000)28    # return pred[0]["generated_text"]29    pred = llama(prompt, temperature=0.1, max_tokens=1000)30    return pred["choices"][0]["text"]31 32def get_first_parquet(dataset: str):33    resp = client.get(f"{BASE_DATASETS_SERVER_URL}/parquet?dataset={dataset}")34    return resp.json()["parquet_files"][0]35 36 37def text2sql(dataset_name, query_input):38    print(f"start text2sql for {dataset_name}")39    try:40        first_parquet = get_first_parquet(dataset_name)41    except Exception as error:42        return {43            schema_output: "",44            prompt_output: "",45            query_output: "",46            df:pd.DataFrame([{"error": f"❌ Could not get dataset schema. {error=}"}])47        }48 49    first_parquet_url = first_parquet["url"]50    print(f"getting schema from {first_parquet_url}")51    con = duckdb.connect()52    con.execute("INSTALL 'httpfs'; LOAD httpfs;")53    # could get from Parquet instead?54    con.execute(f"CREATE TABLE data as SELECT * FROM '{first_parquet_url}' LIMIT 1;")55    result = con.sql("SELECT sql FROM duckdb_tables() where table_name ='data';").df()56    ddl_create = result.iloc[0,0]57    58    text = f"""### Instruction:59    Your task is to generate valid duckdb SQL to answer the following question.60 61    ### Input:62    Here is the database schema that the SQL query will run on:63    {ddl_create}64    65    ### Question:66    {query_input}67 68    ### Response (use duckdb shorthand if possible):69    """70    try:71        sql_output = generate_sql(text)72    except Exception as error:73        return {74            schema_output: ddl_create,75            prompt_output: text,76            query_output: "",77            df:pd.DataFrame([{"error": f"❌ Unable to get the SQL query based on the text. {error=}"}])78        }79 80    # Should be replaced by the prompt but not working81    sql_output = sql_output.replace("FROM data", f"FROM '{first_parquet_url}'")82    try:83        query_result = con.sql(sql_output).df()84    except Exception as error:85        query_result = pd.DataFrame([{"error": f"❌ Could not execute SQL query {error=}"}])86    finally:87        con.close()88    return {89        schema_output: ddl_create,90        prompt_output: text,91        query_output:sql_output,92        df:query_result93    }94 95 96with gr.Blocks() as demo:97    gr.Markdown("# 💫 Generate SQL queries based on a given text for your Hugging Face Dataset 💫")98    dataset_name = HuggingfaceHubSearch(99            label="Hub Dataset ID",100            placeholder="Search for dataset id on Huggingface",101            search_type="dataset",102            value="jamescalam/world-cities-geo",103        )104    # dataset_name = gr.Textbox("jamescalam/world-cities-geo", label="Dataset Name")105    query_input = gr.Textbox("Cities from Albania country", label="Ask something about your data")106    examples = [107                ["Cities from Albania country"],108                ["The continent with the most number of countries"],109                ["Cities that start with 'A'"],110                ["Cities by region"],111            ]112    gr.Examples(examples=examples, inputs=[query_input],outputs=[])113    btn = gr.Button("Generate SQL")114    query_output = gr.Textbox(label="Output SQL", interactive= False)    115    df = gr.DataFrame(datatype="markdown")116    with gr.Accordion("Open for prompt details", open=False):117        #with gr.Column(scale=1, min_width=600):118        schema_output = gr.Textbox(label="Parquet Schema as CREATE DDL", interactive= False)119        prompt_output = gr.Textbox(label="Generated prompt", interactive= False)120    btn.click(text2sql, inputs=[dataset_name, query_input], outputs=[schema_output, prompt_output, query_output,df])121demo.launch(debug=True)122