Update app.py

app.py

@@ -1,53 +1,29 @@
-import json
 import os
-import urllib.parse
+import json
 import gradio as gr
-import requests
-from huggingface_hub import InferenceClient
-
-# InferenceClient setup (you must add your HF token as an environment variable in Colab)
-client = InferenceClient(
-    "meta-llama/Meta-Llama-3.1-70B-Instruct",
-    token=os.environ["HF_TOKEN"],
-)
+import duckdb
 
-# Function to generate iframe for dataset viewer
-def get_iframe(hub_repo_id, sql_query=None):
-    if not hub_repo_id:
-        raise ValueError("Hub repo id is required")
-    if sql_query:
-        sql_query = urllib.parse.quote(sql_query)
-        url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer?sql_console=true&sql={sql_query}"
-    else:
-        url = f"https://huggingface.co/datasets/{hub_repo_id}/embed/viewer"
-    iframe = f"""
-    <iframe src="{url}" frameborder="0" width="100%" height="800px"></iframe>
-    """
-    return iframe
+# Load the Parquet dataset
+dataset_path = '/content/sample_contract_df.parquet'  # Update with your Parquet file's path
 
-# Function to fetch dataset column information
-def get_column_info(hub_repo_id):
-    url = f"https://datasets-server.huggingface.co/info?dataset={hub_repo_id}"
-    response = requests.get(url)
-    try:
-        data = response.json()
-        dataset_info = data.get("dataset_info")
-        key = list(dataset_info.keys())[0]
-        features = json.dumps(dataset_info.get(key).get("features"), indent=2)
-    except Exception as e:
-        return f"Error getting column info: {e}"
-    return features
+# Load the dataset with DuckDB and get schema information
+def load_dataset():
+    con = duckdb.connect()
+    con.execute(f"CREATE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
+    schema = con.execute("DESCRIBE contract_data").fetchdf()
+    con.close()
+    return schema.to_dict(orient="records")
 
-# Function to generate SQL query based on natural language input
-def query_dataset(hub_repo_id, features, query):
+# Generate SQL based on schema and user query
+def generate_sql_query(features, query):
     messages = [
         {
             "role": "system",
-            "content": "You are a SQL query expert assistant that returns a DuckDB SQL query based on the user's natural language query and dataset features.",
+            "content": "You are a SQL query expert assistant that generates DuckDB SQL queries based on the user's natural language query and dataset schema.",
         },
         {
             "role": "user",
-            "content": f"""table train
+            "content": f"""table contract_data
 # Features
 {features}
 # Query
@@ -55,36 +31,53 @@ def query_dataset(hub_repo_id, features, query):
 """,
         },
     ]
-    response = client.chat_completion(
-        messages=messages,
-        max_tokens=1000,
-        stream=False,
-    )
-    query = response.choices[0].message.content
-    return query, get_iframe(hub_repo_id, query)
+    # Here we use DuckDB directly instead of an external API
+    sql_query = f"SELECT * FROM contract_data WHERE {query}"  # Simple example; adapt for complex queries
+    return sql_query
+
+# Execute the SQL query and display results
+def execute_query(sql_query):
+    con = duckdb.connect()
+    con.execute(f"CREATE VIEW contract_data AS SELECT * FROM '{dataset_path}'")
+    result_df = con.execute(sql_query).fetchdf()
+    con.close()
+    return result_df.to_markdown()  # Convert result to markdown for display
 
 # Gradio app UI
 with gr.Blocks() as demo:
     gr.Markdown("""
-    # 🐥 🦙 🤗 Text To SQL Hub Datasets 🤗 🦙 🐥
-    Use this tool to search and query datasets on Huggingface Hub.
-    Built with DuckDB, Huggingface's Inference API, and LLaMA 3.1 70B.
+    # Local Parquet SQL Query App
+    Query and explore the data stored in `sample_contract_df.parquet` using DuckDB and SQL queries.
     """)
+    
+    # Display schema
+    schema = load_dataset()
+    features = json.dumps(schema, indent=2)
+    gr.Markdown(f"### Dataset Schema:\n\n```json\n{features}\n```")
+
+    # User inputs for natural language query
+    query = gr.Textbox(label="Natural Language Query", placeholder="Enter a condition, e.g., 'amount > 1000'")
+    sql_out = gr.Code(label="Generated SQL Query", language="sql")
+    results_out = gr.Markdown(label="Query Results")
+
+    # Buttons to generate and execute SQL
     with gr.Row():
-        with gr.Column():
-            search_in = gr.Textbox(label="Search Huggingface Hub", placeholder="Search for datasets")
-            query = gr.Textbox(label="Natural Language Query", placeholder="Enter a query to generate SQL")
-            sql_out = gr.Code(label="SQL Query", language="sql")
-    with gr.Row():
-        btn = gr.Button("Show Dataset")
-        btn2 = gr.Button("Query Dataset")
-    with gr.Row():
-        search_out = gr.HTML(label="Search Results")
-        features = gr.Code(label="Features", language="json")
+        btn_generate = gr.Button("Generate SQL")
+        btn_execute = gr.Button("Execute Query")
 
-    # Event handling
-    btn.click(fn=get_iframe, inputs=[search_in], outputs=[search_out])
-    btn2.click(fn=query_dataset, inputs=[search_in, features, query], outputs=[sql_out, search_out])
+    # Generate SQL on button click
+    btn_generate.click(
+        fn=generate_sql_query,
+        inputs=[features, query],
+        outputs=sql_out,
+    )
+
+    # Execute SQL on button click
+    btn_execute.click(
+        fn=execute_query,
+        inputs=sql_out,
+        outputs=results_out,
+    )
 
 # Launch the app
 demo.launch()