add llm for generating sql

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+. filter=lfs diff=lfs merge=lfs -text
+Hermes-2-Pro-Llama-3-8B-Q8_0.gguf filter=lfs diff=lfs merge=lfs -text

Hermes-2-Pro-Llama-3-8B-Q8_0.gguf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d138388cfda04d185a68eaf2396cf7a5cfa87d038a20896817a9b7cf1806f532
+size 8541050176

app.py

@@ -4,16 +4,36 @@ 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,
+)
+
 
-def get_dataset_ddl(dataset_id: str) -> pd.DataFrame:
-    view_name = "dataset_view"
+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
 
@@ -43,24 +63,61 @@ CREATE TABLE {} (
     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_name = HuggingfaceHubSearch(
+    dataset_id = HuggingfaceHubSearch(
         label="Hub Dataset ID",
         placeholder="Find your favorite dataset...",
         search_type="dataset",
         value="jamescalam/world-cities-geo",
     )
-    query_input = gr.Textbox("", label="Ask anything...")
+    user_query = gr.Textbox("", label="Ask anything...")
 
     btn = gr.Button("Ask 🪄")
-    df = gr.DataFrame(datatype="markdown")
-    ddl = gr.Text("")
+
+    df = gr.DataFrame()
+    sql_query = gr.Markdown(label="Output SQL Query")
 
     btn.click(
-        get_dataset_ddl,
-        inputs=[dataset_name],
-        outputs=[ddl],
+        query_dataset,
+        inputs=[dataset_id, user_query],
+        outputs=[df, sql_query],
     )
 
 

requirements.txt

@@ -1,7 +1,10 @@
 aiofiles==23.2.1
+aiohttp==3.9.5
+aiosignal==1.3.1
 altair==5.3.0
 annotated-types==0.7.0
 anyio==4.4.0
+async-timeout==4.0.3
 attrs==23.2.0
 certifi==2024.6.2
 charset-normalizer==3.3.2
@@ -9,7 +12,9 @@ click==8.1.7
 contourpy==1.2.1
 cycler==0.12.1
 diskcache==5.6.3
+distro==1.9.0
 dnspython==2.6.1
+docstring_parser==0.16
 duckdb==1.0.0
 email_validator==2.1.1
 exceptiongroup==1.2.1
@@ -18,6 +23,7 @@ fastapi-cli==0.0.4
 ffmpy==0.3.2
 filelock==3.14.0
 fonttools==4.53.0
+frozenlist==1.4.1
 fsspec==2024.6.0
 gradio==4.32.2
 gradio_client==0.17.0
@@ -29,6 +35,7 @@ httpx==0.27.0
 huggingface-hub==0.23.2
 idna==3.7
 importlib_resources==6.4.0
+instructor==1.3.2
 Jinja2==3.1.4
 jsonschema==4.22.0
 jsonschema-specifications==2023.12.1
@@ -39,8 +46,10 @@ MarkupSafe==2.1.5
 matplotlib==3.9.0
 mdurl==0.1.2
 mpmath==1.3.0
+multidict==6.0.5
 networkx==3.3
 numpy==1.26.4
+openai==1.31.0
 orjson==3.10.3
 packaging==24.0
 pandas==2.2.2
@@ -68,6 +77,7 @@ sniffio==1.3.1
 spaces==0.28.3
 starlette==0.37.2
 sympy==1.12.1
+tenacity==8.3.0
 tomlkit==0.12.0
 toolz==0.12.1
 torch==2.3.0
@@ -81,3 +91,4 @@ uvicorn==0.30.1
 uvloop==0.19.0
 watchfiles==0.22.0
 websockets==11.0.3
+yarl==1.9.4