Update app.py

app.py

@@ -1,104 +1,93 @@
-import gradio as gr
-import torch
-import pyperclip
-import openai
-import os
-import pandas as pd
-from sqlalchemy import create_engine, inspect
-from llama_index.legacy import (
-    VectorStoreIndex,
-    SQLDatabase,
-    ServiceContext,
-)
-from llama_index.legacy.indices.struct_store import NLSQLTableQueryEngine
-from llama_index.legacy.llms import OpenAI
-import sqlite3
-
-# Set up OpenAI API Key
-# Load the OpenAI API key from the environment variable
-openai.api_key = os.getenv("OPENAI_API_KEY")
+# Install necessary libraries
+!pip install gradio requests duckdb huggingface_hub
 
-# Check if the API key is set
-if not openai.api_key:
-    raise ValueError("OPENAI_API_KEY environment variable is not set")
-
-# Function to load database and LLM
-def load_db_llm():
-    engine = create_engine("sqlite:///gov-contracts.db")
-    sql_database = SQLDatabase(engine)
-    llm = OpenAI(temperature=0.1, model="gpt-3.5-turbo-1106")
-    service_context = ServiceContext.from_defaults(llm=llm, embed_model="local")
-    return sql_database, service_context, engine
+import json
+import os
+import urllib.parse
+import gradio as gr
+import requests
+from huggingface_hub import InferenceClient
 
-# Load LLM and database context
-sql_database, service_context, engine = load_db_llm()
-query_engine = NLSQLTableQueryEngine(
-    sql_database=sql_database, synthesize_response=True, service_context=service_context
+# 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"],
 )
 
-# Initialize table schema and connection for query
-inspector = inspect(engine)
-table_names = inspector.get_table_names()
+# 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 table data function
-def get_table_data(table_name):
-    conn = sqlite3.connect('gov-contracts.db')
-    query = f"SELECT * FROM {table_name}"
-    df = pd.read_sql_query(query, conn)
-    conn.close()
-    return df
+# 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
 
-# Chat-based interaction for Gradio
-def generate_response(user_input, selected_table=None, example_prompt=None):
-    if example_prompt:
-        user_input = example_prompt
-    
-    response = query_engine.query(f"User Question: {user_input}")
-    sql_query = f"```sql\n{response.metadata['sql_query']}\n```\n**Response:**\n{response.response}\n"
-    pyperclip.copy(sql_query)  # Optional: Copy to clipboard
-    return sql_query
+# Function to generate SQL query based on natural language input
+def query_dataset(hub_repo_id, 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.",
+        },
+        {
+            "role": "user",
+            "content": f"""table train
+# Features
+{features}
+# Query
+{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)
 
-# Define Gradio app layout and components
-with gr.Blocks() as gradio_app:
-    gr.Markdown("## Natural Language to SQL Query Application")
-    gr.Markdown("### Ask a question about the data in the database to receive a precise SQL query.")
-    
-    # Sidebar: Database schema and example prompts
+# 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.
+    """)
     with gr.Row():
         with gr.Column():
-            table_dropdown = gr.Dropdown(choices=table_names, label="Select a Table")
-            example_prompt_box = gr.Radio(
-                choices=[
-                    "Return the department_ind_agency and the sum of award in descending order", 
-                    "Return the sum of award in descending order grouped by type limited to the top 10", 
-                    "Return the sum of award by year where the sub_tier is the FEDERAL ACQUISITION SERVICE"
-                ],
-                label="Select an Example Prompt"
-            )
-            query_btn = gr.Button("Generate Query")
-
-    with gr.Column():
-        user_query = gr.Textbox(
-            label="Enter your natural language query about the database",
-            placeholder="Ask your question here..."
-        )
-        chat_output = gr.Textbox(
-            label="Generated SQL Query",
-            placeholder="SQL query will appear here..."
-        )
-
-    # Function to call on click
-    def query_callback(user_input, table_name, example_prompt):
-        return generate_response(user_input, selected_table=table_name, example_prompt=example_prompt)
-
-    # Button click event
-    query_btn.click(query_callback, inputs=[user_query, table_dropdown, example_prompt_box], outputs=chat_output)
+            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")
 
-    gr.Markdown("#### Created by Leonce Nshuti")
-    gr.Markdown("""
-        - [LinkedIn](https://www.linkedin.com/in/leoncenshuti/)
-        - [GitHub](https://github.com/LNshuti)
-    """)
+    # 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])
 
-if __name__ == "__main__":
-    gradio_app.launch()
+# Launch the app
+demo.launch()