Update app.py

app.py

@@ -6,6 +6,7 @@ from httpx import Client
 from huggingface_hub import HfApi
 from huggingface_hub.utils import logging
 from llama_cpp import Llama
+import pandas as pd
 
 load_dotenv()
 
@@ -24,10 +25,13 @@ headers = {
 logger = logging.get_logger(__name__)
 client = Client(headers=headers)
 api = HfApi(token=HF_TOKEN)
+
+print("About to load DuckDB-NSQL-7B model")
 llama = Llama(
         model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
         n_ctx=2048,
 )
+print("DuckDB-NSQL-7B model has been loaded")
 
 def get_first_parquet(dataset: str):
     resp = client.get(f"{BASE_DATASETS_SERVER_URL}/parquet?dataset={dataset}")
@@ -59,9 +63,9 @@ def text2sql(dataset_name, query_input):
     print(first_parquet_url)
     con = duckdb.connect()
     con.execute("INSTALL 'httpfs'; LOAD httpfs;")
+    # could get from parquet instead?
     con.execute(f"CREATE TABLE data as SELECT * FROM '{first_parquet_url}' LIMIT 1;")
     result = con.sql("SELECT sql FROM duckdb_tables() where table_name ='data';").df()
-    con.close()
 
     ddl_create = result.iloc[0,0]
     text = f"""### Instruction:
@@ -73,7 +77,7 @@ def text2sql(dataset_name, query_input):
     ### Question:
     {query_input}
 
-    ### Response (use duckdb shorthand if possible):
+    ### Response (use duckdb shorthand if possible) replace table name with {first_parquet_url} in the generated sql query:
     """
     
     print(text)
@@ -81,7 +85,17 @@ def text2sql(dataset_name, query_input):
     # sql_output =  query_remote_model(text)
 
     sql_output = query_local_model(text)
-    return sql_output
+
+    try:
+        query_result = con.sql(sql_output).df()
+    except Exception as error:
+        query_result = pd.DataFrame([{"error": f"❌ Could not execute SQL query {error=}"}])
+    finally:
+        con.close()
+    return {
+        query_output:sql_output,
+        df:query_result
+    }
 
 
 with gr.Blocks() as demo:
@@ -92,5 +106,6 @@ with gr.Blocks() as demo:
     query_input = gr.Textbox("How many rows there are?", label="Ask something about your data")
     btn = gr.Button("Generate SQL")
     query_output = gr.Textbox(label="Output SQL", interactive= False)    
-    btn.click(text2sql, inputs=[dataset_name, query_input], outputs=query_output)
+    df = gr.DataFrame(datatype="markdown")
+    btn.click(text2sql, inputs=[dataset_name, query_input], outputs=[query_output,df])
 demo.launch()