app.py

"""
TODOS:
- Improve prompts
- Improve model usage (Quantization?)
- Improve error handling
- Add more tests
- Improve response in a friendly way
"""

import gradio as gr
from gradio_huggingfacehub_search import HuggingfaceHubSearch
import duckdb
import pandas as pd
import requests
from outlines import prompt
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces
import json
import torch
import logging

BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
logger = logging.getLogger(__name__)

"""
Methods for generating potential questions and SQL queries
"""
device = "cuda"
gemma_model_id = "google/gemma-2b-it"
gemma_tokenizer = AutoTokenizer.from_pretrained(gemma_model_id)
gemma_model = AutoModelForCausalLM.from_pretrained(
    gemma_model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16
)

@spaces.GPU
def generate_potential_questions_with_gemma(prompt):
    input_ids = gemma_tokenizer(prompt, return_tensors="pt").to(device)
    outputs = gemma_model.generate(**input_ids, max_new_tokens=1024)
    return gemma_tokenizer.decode(outputs[0], skip_special_tokens=True)


@prompt
def prompt_for_questions(dataset, schema, first_rows):
    """
    You are a data analyst tasked with exploring a dataset named {{ dataset }}.
    Below is the dataset schema in SQL format along with a sample of 3 rows:
    {{ schema }}
    Sample rows:
    {% for example in first_rows %}
    {{ example}}
    {% endfor %}
    Your goal is to generate a list of 5 potential questions that a user might want
    to ask about this dataset. Consider the information contained in the provided
    columns and rows, and try to think of meaningful questions that could
    provide insights or useful information. For each question, provide the SQL query
    that would extract the relevant information from the dataset.
    Ouput JSON format:
    {
        "questions": [
            {"question": [Insert question here]", "sql_query": "[Insert SQL query here]"},
            {"question": [Insert question here]", "sql_query": "[Insert SQL query here]"},
            {"question": [Insert question here]", "sql_query": "[Insert SQL query here]"},
            {"question": [Insert question here]", "sql_query": "[Insert SQL query here]"},
            {"question": [Insert question here]", "sql_query": "[Insert SQL query here]"},
        ]
    }
    Please ensure that each SQL query retrieves relevant information from the dataset to answer the corresponding question accurately.
    Return only the JSON object, do not add extra information.
    """

"""
Methods for generating and SQL based on a user request 
"""
mother_duckdb_model_id = "motherduckdb/DuckDB-NSQL-7B-v0.1"
mother_duck_tokenizer = AutoTokenizer.from_pretrained(mother_duckdb_model_id)
mother_duck_model = AutoModelForCausalLM.from_pretrained(
    mother_duckdb_model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16
)

@spaces.GPU
def generate_sql_with_mother_duck(prompt):
    input_ids = mother_duck_tokenizer(prompt, return_tensors="pt").to(device).input_ids
    generated_ids = mother_duck_model.generate(input_ids, max_length=1024)
    return mother_duck_tokenizer.decode(generated_ids[0], skip_special_tokens=True)


@prompt
def prompt_for_sql(ddl_create, query_input):
    """
    ### Instruction:
    Your task is to generate valid duckdb SQL to answer the following question.
    ### Input:
    Here is the database schema that the SQL query will run on:
    {{ ddl_create }}
    
    ### Question:
    {{ query_input }}
    ### Response (use duckdb shorthand if possible):
    """


"""
Datasets Viewer Methods
https://huggingface.co/docs/datasets-server/index
"""

def get_first_parquet(dataset: str):
    resp = requests.get(f"{BASE_DATASETS_SERVER_URL}/parquet?dataset={dataset}")
    return resp.json()["parquet_files"][0]


def get_dataset_schema(parquet_url: str):
    con = duckdb.connect()
    con.execute(f"CREATE TABLE data as SELECT * FROM '{parquet_url}' LIMIT 1;")
    result = con.sql("SELECT sql FROM duckdb_tables() where table_name ='data';").df()
    ddl_create = result.iloc[0,0]
    con.close()
    return ddl_create


def get_first_rows_as_df(dataset: str, config: str, split: str, limit:int):
    resp = requests.get(f"{BASE_DATASETS_SERVER_URL}/first-rows?dataset={dataset}&config={config}&split={split}")
    rows = resp.json()["rows"]
    rows = [row['row'] for row in rows]
    return pd.DataFrame.from_dict(rows).sample(frac = 1).head(limit)

"""
Main logic, to get the recommended queries
"""
def get_recommended_queries(dataset: str):
    ddl_create, prompt = "", ""
    try:
        first_split = get_first_parquet(dataset)
        df_first_rows = get_first_rows_as_df(dataset, first_split["config"], first_split["split"], 3)
        first_parquet_url = first_split["url"]
        logger.info(f"First parquet URL: {first_parquet_url}")
        ddl_create = get_dataset_schema(first_parquet_url)
        prompt = prompt_for_questions(dataset, ddl_create, df_first_rows.to_dict('records'))
        txt_questions = generate_potential_questions_with_gemma(prompt).split("``json")[1].replace('\n', ' ').strip()[:-4]
        data = json.loads(txt_questions)
        questions = data["questions"]
        potential_questions = []
        for question in questions:
            try:
                sql = question["sql_query"].replace("FROM data", f"FROM '{first_parquet_url}'")
                result = duckdb.sql(sql).df()
                potential_questions.append({"question": question["question"], "result": result, "sql_query": sql})
                continue
            except Exception as err:
                logger.error(f"Error in running SQL query: {question['sql_query']} {err}")
                mother_duck_prompt = prompt_for_sql(ddl_create, question["question"])
                sql = generate_sql_with_mother_duck(mother_duck_prompt).split("### Response (use duckdb shorthand if possible):")[-1].strip()
                sql = sql.replace("FROM data", f"FROM '{first_parquet_url}'")
                try:
                    result = duckdb.sql(sql).df()
                    potential_questions.append({"question": question["question"], "result": result, "sql_query": sql})
                except:
                    pass
        df_result = pd.DataFrame(potential_questions)
    except Exception as err:
        logger.error(f"Error in getting recommended queries: {err}")
        return {
            gr_txt_ddl: ddl_create,
            gr_txt_prompt: prompt,
            gr_df_result: pd.DataFrame([{"error": f"❌ {err=}"}])
        }
    return {
        gr_txt_ddl: ddl_create,
        gr_txt_prompt: prompt,
        gr_df_result: df_result
    }


def preview_dataset(dataset: str):
    try:
        first_split = get_first_parquet(dataset)
        df = get_first_rows_as_df(dataset, first_split["config"], first_split["split"], 4)
    except Exception as err:
        df = pd.DataFrame([{"Unable to preview dataset": f"❌ {err=}"}])
    return {
        gr_df_first_rows: df
    }


with gr.Blocks() as demo:
    gr.Markdown("# 💫 Dataset Insights Explorer 💫")
    gr_dataset_name = HuggingfaceHubSearch(
            label="Hub Dataset ID",
            placeholder="Search for dataset id on Huggingface",
            search_type="dataset",
            value="jamescalam/world-cities-geo",
        )
    gr_preview_btn = gr.Button("Preview Dataset")
    gr_df_first_rows = gr.DataFrame(datatype="markdown")
    gr_recommend_btn = gr.Button("Show Insights")
    gr_df_result = gr.DataFrame(datatype="markdown")
    with gr.Accordion("Open for details", open=False):
        gr_txt_ddl = gr.Textbox(label="Dataset as CREATE DDL", interactive= False)
        gr_txt_prompt = gr.Textbox(label="Generated prompt to get recommended questions", interactive= False)
    gr_preview_btn.click(preview_dataset, inputs=[gr_dataset_name], outputs=[gr_df_first_rows])
    gr_recommend_btn.click(get_recommended_queries, inputs=[gr_dataset_name], outputs=[gr_txt_ddl, gr_txt_prompt, gr_df_result])
demo.launch()