Update app.py

app.py

@@ -1,19 +1,52 @@
-import gradio as gr
+import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import re
 import sqlparse
+import time
 
-# Load model and tokenizer
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForCausalLM.from_pretrained(
-    "onkolahmet/Qwen2-0.5B-Instruct-SQL-generator", 
-    torch_dtype="auto",
-    device_map="auto"
-)
-tokenizer = AutoTokenizer.from_pretrained("onkolahmet/Qwen2-0.5B-Instruct-SQL-generator")
+# App title and description
+st.set_page_config(page_title="Text to SQL Converter", layout="wide")
+st.title("Text to SQL Query Converter")
+st.markdown("Enter your question and optional table context to generate an SQL query.")
 
-def generate_sql(question, context=None):
+# Model loading (with loading state indicator)
+@st.cache_resource
+def load_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    st.info(f"Loading model on {device}... This may take a minute.")
+    
+    # Load without device_map for compatibility
+    model = AutoModelForCausalLM.from_pretrained(
+        "onkolahmet/text_to_sql", 
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    model = model.to(device)
+    tokenizer = AutoTokenizer.from_pretrained("onkolahmet/text_to_sql")
+    
+    return model, tokenizer, device
+
+# Few-shot examples to include in each prompt
+examples = [
+    {
+        "question": "Get the names and emails of customers who placed an order in the last 30 days.",
+        "sql": "SELECT name, email FROM customers WHERE order_date >= DATE_SUB(CURDATE(), INTERVAL 30 DAY);"
+    },
+    {
+        "question": "Find all employees with a salary greater than 50000.",
+        "sql": "SELECT * FROM employees WHERE salary > 50000;"
+    },
+    {
+        "question": "List all product names and their categories where the price is below 50.",
+        "sql": "SELECT name, category FROM products WHERE price < 50;"
+    },
+    {
+        "question": "How many users registered in the year 2022?",
+        "sql": "SELECT COUNT(*) FROM users WHERE YEAR(registration_date) = 2022;"
+    }
+]
+
+def generate_sql(model, tokenizer, device, question, context=None):
     # Construct prompt with few-shot examples and context if available
     prompt = "Translate natural language questions to SQL queries.\n\n"
     
@@ -21,6 +54,9 @@ def generate_sql(question, context=None):
     if context and context.strip():
         prompt += f"Table Context:\n{context}\n\n"
     
+    # Add few-shot examples
+    for ex in examples:
+        prompt += f"Q: {ex['question']}\nSQL: {ex['sql']}\n\n"
     
     # Add the current question
     prompt += f"Q: {question}\nSQL:"
@@ -146,36 +182,39 @@ def clean_sql_output(sql_text):
     except:
         return best_query
 
-def process_input(question, table_context):
-    """Function to process user input through the model and return formatted results"""
-    if not question.strip():
-        return "Please enter a question."
-    
-    # Generate SQL from the question and context
-    raw_sql = generate_sql(question, table_context)
-    
-    # Clean the SQL output
-    cleaned_sql = clean_sql_output(raw_sql)
-    
-    if not cleaned_sql:
-        return "Sorry, I couldn't generate a valid SQL query. Please try rephrasing your question."
-    
-    return cleaned_sql
+# Load model (this happens once when the app starts)
+model, tokenizer, device = load_model()
+st.success("Model loaded successfully! Ready to generate SQL queries.")
 
-# Sample table context examples for the example selector
-example_contexts = [
-    # Example 1
-    """
-CREATE TABLE customers (
-  id INT PRIMARY KEY,
-  name VARCHAR(100),
-  email VARCHAR(100),
-  order_date DATE
-);
-    """,
+# Main app interface
+col1, col2 = st.columns([1, 1])
+
+with col1:
+    # User inputs
+    question = st.text_area("Your Question", 
+                            placeholder="e.g., Find all products with price less than $50",
+                            height=100)
     
-    # Example 2
-    """
+    table_context = st.text_area("Table Context (Optional)", 
+                                placeholder="Enter your database schema or table definitions here...", 
+                                height=200)
+    
+    # Example selection
+    with st.expander("Try an example", expanded=False):
+        example_option = st.selectbox(
+            "Select an example:",
+            [
+                "List all products in the 'Electronics' category with price less than $500",
+                "Find the total number of employees in each department",
+                "Get customers who placed orders in the last 7 days",
+                "Count the number of products in each category",
+                "Find the average salary by department"
+            ]
+        )
+        
+        # Sample table context examples mapped to questions
+        example_contexts = {
+            "List all products in the 'Electronics' category with price less than $500": """
 CREATE TABLE products (
   id INT PRIMARY KEY,
   name VARCHAR(100),
@@ -183,10 +222,8 @@ CREATE TABLE products (
   price DECIMAL(10,2),
   stock_quantity INT
 );
-    """,
-    
-    # Example 3
-    """
+            """,
+            "Find the total number of employees in each department": """
 CREATE TABLE employees (
   id INT PRIMARY KEY,
   name VARCHAR(100),
@@ -201,87 +238,97 @@ CREATE TABLE departments (
   manager_id INT,
   budget DECIMAL(15,2)
 );
-    """
-]
+            """,
+            "Get customers who placed orders in the last 7 days": """
+CREATE TABLE customers (
+  id INT PRIMARY KEY,
+  name VARCHAR(100),
+  email VARCHAR(100),
+  order_date DATE
+);
+            """,
+            "Count the number of products in each category": """
+CREATE TABLE products (
+  id INT PRIMARY KEY,
+  name VARCHAR(100),
+  category VARCHAR(50),
+  price DECIMAL(10,2),
+  stock_quantity INT
+);
+            """,
+            "Find the average salary by department": """
+CREATE TABLE employees (
+  id INT PRIMARY KEY,
+  name VARCHAR(100),
+  department VARCHAR(50),
+  salary DECIMAL(10,2),
+  hire_date DATE
+);
+            """
+        }
+        
+        apply_example = st.button("Apply Example")
+        if apply_example:
+            question = example_option
+            table_context = example_contexts[example_option]
+            st.session_state.question = question
+            st.session_state.table_context = table_context
+            st.success("Example applied! Click 'Generate SQL Query' to see the result.")
 
-# Sample question examples
-example_questions = [
-    "Get the names and emails of customers who placed an order in the last 30 days.",
-    "Find all products with less than 10 items in stock.",
-    "List all employees in the Sales department with a salary greater than 50000.",
-    "What is the total budget for departments with more than 5 employees?",
-    "Count how many products are in each category where the price is greater than 100."
-]
+    # Button to generate SQL
+    generate_button = st.button("Generate SQL Query")
 
-# Create the Gradio interface
-with gr.Blocks(title="Text to SQL Converter") as demo:
-    gr.Markdown("# Text to SQL Query Converter")
-    gr.Markdown("Enter your question and optional table context to generate an SQL query.")
-    
-    with gr.Row():
-        with gr.Column():
-            question_input = gr.Textbox(
-                label="Your Question",
-                placeholder="e.g., Find all products with price less than $50",
-                lines=2
-            )
+# Display results
+with col2:
+    if generate_button and question:
+        with st.spinner("Generating SQL query..."):
+            # Record start time
+            start_time = time.time()
             
-            table_context = gr.Textbox(
-                label="Table Context (Optional)",
-                placeholder="Enter your database schema or table definitions here...",
-                lines=10
-            )
+            # Generate SQL
+            raw_sql = generate_sql(model, tokenizer, device, question, table_context)
+            cleaned_sql = clean_sql_output(raw_sql)
             
-            submit_btn = gr.Button("Generate SQL Query")
-        
-        with gr.Column():
-            sql_output = gr.Code(
-                label="Generated SQL Query",
-                language="sql",
-                lines=12
+            # Calculate elapsed time
+            elapsed_time = time.time() - start_time
+            
+            # Display results
+            st.subheader("Generated SQL Query")
+            st.code(cleaned_sql, language="sql")
+            st.info(f"Query generated in {elapsed_time:.2f} seconds")
+            
+            # Display explanation
+            st.subheader("Explanation")
+            st.write("This SQL query translates your natural language question into a database command.")
+            
+            # Option to copy to clipboard (using JavaScript)
+            st.markdown(
+                f"""
+                <div style="margin-top: 20px;">
+                    <button 
+                        onclick="navigator.clipboard.writeText(`{cleaned_sql}`);this.textContent='Copied!';setTimeout(()=>this.textContent='Copy to Clipboard',1500)"
+                        style="background-color:#4CAF50;color:white;padding:8px 16px;border:none;border-radius:4px;cursor:pointer"
+                    >
+                        Copy to Clipboard
+                    </button>
+                </div>
+                """,
+                unsafe_allow_html=True
             )
-    
-    # Examples section
-    gr.Markdown("### Try some examples")
-    
-    example_selector = gr.Examples(
-        examples=[
-            ["List all products in the 'Electronics' category with price less than $500", example_contexts[1]],
-            ["Find the total number of employees in each department", example_contexts[2]],
-            ["Get customers who placed orders in the last 7 days", example_contexts[0]],
-            ["Count the number of products in each category", example_contexts[1]],
-            ["Find the average salary by department", example_contexts[2]]
-        ],
-        inputs=[question_input, table_context]
-    )
-    
-    # Set up the submit button to trigger the process_input function
-    submit_btn.click(
-        fn=process_input,
-        inputs=[question_input, table_context],
-        outputs=sql_output
-    )
-    
-    # Also trigger on pressing Enter in the question input
-    question_input.submit(
-        fn=process_input,
-        inputs=[question_input, table_context],
-        outputs=sql_output
-    )
-    
-    # Add information about the model
-    gr.Markdown("""
-    ### About
-    This app uses a fine-tuned language model to convert natural language questions into SQL queries.
-    
-    - **Model**: [onkolahmet/Qwen2-0.5B-Instruct-SQL-generator](https://huggingface.co/onkolahmet/Qwen2-0.5B-Instruct-SQL-generator)
-    - **How to use**: 
-      1. Enter your question in natural language
-      2. If you have specific table schemas, add them in the Table Context field
-      3. Click "Generate SQL Query" or press Enter
-      
-    Note: The model works best when table context is provided, but can generate generic SQL queries without it.
-    """)
+    else:
+        st.info("Enter a question and click 'Generate SQL Query' to see the result here.")
+
+# App footer and info
+st.markdown("---")
+st.markdown("""
+### About
+This app uses a fine-tuned language model to convert natural language questions into SQL queries.
 
-# Launch the app
-demo.launch()
+- **Model**: [onkolahmet/Qwen2-0.5B-Instruct-SQL-generator](https://huggingface.co/onkolahmet/Qwen2-0.5B-Instruct-SQL-generator)
+- **How to use**: 
+  1. Enter your question in natural language
+  2. If you have specific table schemas, add them in the Table Context field
+  3. Click "Generate SQL Query"
+  
+Note: The model works best when table context is provided, but can generate generic SQL queries without it.
+""")