app.py

from gradio_huggingfacehub_search import HuggingfaceHubSearch
from huggingface_hub import HfApi
import pandas as pd
import gradio as gr
import duckdb
import requests
import llama_cpp
import instructor

from pydantic import BaseModel

BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
view_name = "dataset_view"

hf_api = HfApi()
conn = duckdb.connect()

llama = llama_cpp.Llama(
    model_path="Hermes-2-Pro-Llama-3-8B-Q8_0.gguf",
    n_gpu_layers=-1,
    chat_format="chatml",
    n_ctx=2048,
    verbose=False,
)

create = instructor.patch(
    create=llama.create_chat_completion_openai_v1,
    mode=instructor.Mode.JSON_SCHEMA,
)


class SQLResponse(BaseModel):
    sql: str


def get_dataset_ddl(dataset_id: str) -> str:
    response = requests.get(f"{BASE_DATASETS_SERVER_URL}/parquet?dataset={dataset_id}")
    response.raise_for_status()  # Check if the request was successful

    first_parquet = response.json().get("parquet_files", [])[0]
    first_parquet_url = first_parquet.get("url")

    if not first_parquet_url:
        raise ValueError("No valid URL found for the first parquet file.")

    conn.execute(
        f"CREATE OR REPLACE VIEW {view_name} as SELECT * FROM read_parquet('{first_parquet_url}');"
    )
    dataset_ddl = conn.execute(f"PRAGMA table_info('{view_name}');").fetchall()

    column_data_types = ",\n\t".join(
        [f"{column[1]} {column[2]}" for column in dataset_ddl]
    )

    sql_ddl = """
CREATE TABLE {} (
    {}
);
    """.format(
        view_name, column_data_types
    )

    return sql_ddl


def generate_sql(dataset_id: str, query: str) -> str:
    ddl = get_dataset_ddl(dataset_id)

    system_prompt = f"""
    You are an expert SQL assistant with access to the following DuckDB Table:
    
    ```sql
    {ddl}
    ```
    
    Please assist the user by writing a SQL query that answers the user's question.
    """

    resp: SQLResponse = create(
        model="Hermes-2-Pro-Llama-3-8B",
        messages=[
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": query,
            },
        ],
        response_model=SQLResponse,
    )

    return resp.sql


def query_dataset(dataset_id: str, query: str) -> tuple[pd.DataFrame, str]:
    sql_query = generate_sql(dataset_id, query)
    df = conn.execute(sql_query).fetchdf()

    markdown_output = f"""```sql\n{sql_query}```"""
    return df, markdown_output


with gr.Blocks() as demo:
    gr.Markdown("# Query your HF Datasets with Natural Language 📈📊")
    dataset_id = HuggingfaceHubSearch(
        label="Hub Dataset ID",
        placeholder="Find your favorite dataset...",
        search_type="dataset",
        value="jamescalam/world-cities-geo",
    )
    user_query = gr.Textbox("", label="Ask anything...")

    btn = gr.Button("Ask 🪄")

    df = gr.DataFrame()
    sql_query = gr.Markdown(label="Output SQL Query")

    btn.click(
        query_dataset,
        inputs=[dataset_id, user_query],
        outputs=[df, sql_query],
    )


if __name__ == "__main__":
    demo.launch()