Upload 3 files

Add basic requirements

requirements.txt

@@ -0,0 +1,4 @@
+gradio==5.23.1
+peft==0.15.1
+torch==2.6.0+cu118
+transformers==4.50.3

schema_extractor.py

@@ -0,0 +1,132 @@
+import sqlite3
+from typing import Optional
+
+
+class SQLiteSchemaExtractor:
+    """
+    A class to extract the schema from an SQLite database file and return it as formatted text.
+    """
+
+    def __init__(self, db_path: str):
+        """
+        Initialize the extractor with the path to the SQLite database file.
+
+        Args:
+            db_path: Path to the SQLite database file
+        """
+        self.db_path = db_path
+        self.connection = None
+
+    def connect(self):
+        """Establish a connection to the SQLite database."""
+        try:
+            self.connection = sqlite3.connect(self.db_path)
+            return True
+        except sqlite3.Error as e:
+            print(f"Error connecting to database: {e}")
+            return False
+
+    def close(self):
+        """Close the database connection if it exists."""
+        if self.connection:
+            self.connection.close()
+
+    def get_schema(self) -> Optional[str]:
+        """
+        Extract the schema from the SQLite database and return it as formatted text.
+
+        Returns:
+            A string containing the formatted schema, or None if an error occurred
+        """
+        if not self.connect():
+            return None
+
+        try:
+            cursor = self.connection.cursor()
+
+            # Get the list of all tables
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+            tables = cursor.fetchall()
+
+            schema_text = ""
+
+            for table in tables:
+                table_name = table[0]
+
+                # Skip SQLite internal tables
+                if table_name.startswith("sqlite_"):
+                    continue
+
+                # Get the CREATE TABLE statement for the current table
+                cursor.execute(f"PRAGMA table_info({table_name});")
+                columns = cursor.fetchall()
+
+                # Format the CREATE TABLE statement
+                schema_text += f"CREATE TABLE {table_name} (\n"
+
+                for i, col in enumerate(columns):
+                    # Column format: column_name data_type
+                    col_id, col_name, col_type, not_null, default_val, pk = col
+
+                    # Map SQLite types to our simplified type system
+                    simplified_type = self._simplify_type(col_type.lower())
+
+                    # Add column definition
+                    schema_text += f"    {col_name} {simplified_type}"
+
+                    # Add comma if not the last column
+                    if i < len(columns) - 1:
+                        schema_text += ","
+
+                    schema_text += "\n"
+
+                schema_text += ")\n\n"
+
+            # Remove the last newline
+            if schema_text.endswith("\n\n"):
+                schema_text = schema_text[:-1]
+
+            return schema_text
+
+        except sqlite3.Error as e:
+            print(f"Error extracting schema: {e}")
+            return None
+        finally:
+            self.close()
+
+    def _simplify_type(self, sqlite_type: str) -> str:
+        """
+        Convert SQLite types to simplified types (number, text, others).
+
+        Args:
+            sqlite_type: The SQLite data type
+
+        Returns:
+            A simplified type name
+        """
+        if any(
+            num_type in sqlite_type
+            for num_type in ["int", "real", "floa", "doub", "num", "dec"]
+        ):
+            return "number"
+        elif any(
+            text_type in sqlite_type
+            for text_type in ["text", "char", "clob", "varchar"]
+        ):
+            return "text"
+        else:
+            return "others"
+
+
+# # Example usage:
+# if __name__ == "__main__":
+#     # Replace with your SQLite database file path
+#     db_file = "path/to/your/database.db"
+
+#     extractor = SQLiteSchemaExtractor(db_file)
+#     schema = extractor.get_schema()
+
+#     if schema:
+#         print(schema)
+#     else:
+#         print("Failed to extract schema.")

ui.py

@@ -0,0 +1,236 @@
+import os
+import tempfile
+
+import gradio as gr
+import torch
+from peft import PeftConfig, PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from schema_extractor import SQLiteSchemaExtractor
+
+
+# Load model and tokenizer
+def load_model():
+    config = PeftConfig.from_pretrained("Rajan/training_run")
+    tokenizer = AutoTokenizer.from_pretrained("NumbersStation/nsql-350M")
+    base_model = AutoModelForCausalLM.from_pretrained("NumbersStation/nsql-350M")
+    model = PeftModel.from_pretrained(base_model, "Rajan/training_run")
+    return model, tokenizer
+
+
+# Extract and correct SQL from generated text
+def extract_and_correct_sql(text, correct=False):
+    lines = text.splitlines()
+
+    start_index = 0
+    for i, line in enumerate(lines):
+        if line.strip().upper().startswith("SELECT"):
+            start_index = i
+            break
+
+    generated_sql = "\n".join(lines[start_index:])
+
+    if correct:
+        if not generated_sql.strip().endswith(";"):
+            generated_sql = generated_sql.strip() + ";"
+
+    return generated_sql
+
+
+# Function to handle file upload and schema extraction
+def upload_and_extract_schema(db_file):
+    if db_file is None:
+        return "Please upload a database file", None
+
+    try:
+        # Get the file path directly from Gradio
+        temp_db_path = db_file.name
+
+        extractor = SQLiteSchemaExtractor(temp_db_path)
+        schema = extractor.get_schema()
+        return schema, temp_db_path
+    except Exception as e:
+        return f"Error extracting schema: {str(e)}", None
+
+
+# Function to handle chat interaction with streaming effect
+def generate_sql(question, schema, db_path, chat_history):
+    if db_path is None or not schema:
+        return (
+            chat_history
+            + [
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": "Please upload a database file first"},
+            ],
+            None,
+        )
+
+    try:
+        # Load model
+        model, tokenizer = load_model()
+
+        # Format prompt
+        prompt_format = """ 
+{}
+-- Using valid SQLite, answer the following questions for the tables provided above.
+{}
+SELECT"""
+
+        # Format the prompt with schema and question
+        prompt = prompt_format.format(schema, question)
+
+        # Generate SQL
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        generated_ids = model.generate(input_ids, max_length=500)
+        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+        # Extract SQL
+        sql_query = extract_and_correct_sql(generated_text, correct=True)
+
+        # Update history using dictionary format
+        new_history = chat_history + [
+            {"role": "user", "content": question},
+            {"role": "assistant", "content": sql_query},
+        ]
+        return new_history, sql_query
+    except Exception as e:
+        error_message = f"Error: {str(e)}"
+        return (
+            chat_history
+            + [
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": error_message},
+            ],
+            None,
+        )
+
+
+# Function for streaming SQL generation effect
+def stream_sql(question, schema, db_path, chat_history):
+    # First add the user message to chat
+    yield chat_history + [{"role": "user", "content": question}], ""
+
+    if db_path is None or not schema:
+        yield chat_history + [
+            {"role": "user", "content": question},
+            {"role": "assistant", "content": "Please upload a database file first"},
+        ], "Please upload a database file first"
+        return
+
+    try:
+        # Load model
+        model, tokenizer = load_model()
+
+        # Format prompt
+        prompt_format = """ 
+{}
+-- Using valid SQLite, answer the following questions for the tables provided above.
+{}
+SELECT"""
+
+        # Format the prompt with schema and question
+        prompt = prompt_format.format(schema, question)
+
+        # Generate SQL
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        generated_ids = model.generate(input_ids, max_length=500)
+        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+        # Extract SQL
+        sql_query = extract_and_correct_sql(generated_text, correct=True)
+
+        # Fixed medium speed (0.03 seconds delay)
+        import time
+
+        delay = 0.03  # 30ms - normal typing speed
+
+        # Stream the SQL query character by character for effect
+        partial_sql = ""
+        for char in sql_query:
+            partial_sql += char
+            # Update chat history and SQL display with each character
+            yield chat_history + [
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": partial_sql},
+            ], partial_sql
+            time.sleep(delay)  # Medium speed typing effect
+
+    except Exception as e:
+        error_message = f"Error: {str(e)}"
+        yield chat_history + [
+            {"role": "user", "content": question},
+            {"role": "assistant", "content": error_message},
+        ], error_message
+
+
+# Main application
+def create_app():
+    with gr.Blocks(title="SQL Query Generator", theme=gr.themes.Soft()) as app:
+        gr.Markdown("# SQL Query Generator")
+        gr.Markdown(
+            "Upload a SQLite database, ask questions, and get SQL queries automatically generated"
+        )
+
+        # Store database path
+        db_path = gr.State(value=None)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                # File upload section
+                file_input = gr.File(label="Upload SQLite Database (.db file)")
+                extract_btn = gr.Button("Extract Schema", variant="primary")
+
+                # Schema display
+                schema_output = gr.Textbox(
+                    label="Database Schema", lines=10, interactive=False
+                )
+
+            with gr.Column(scale=2):
+                # Chat interface
+                chatbot = gr.Chatbot(
+                    label="Query Conversation", height=400, type="messages"
+                )
+
+                with gr.Row():
+                    question_input = gr.Textbox(
+                        label="Ask a question about your data",
+                        placeholder="e.g., Show me the top 10 most sold items",
+                    )
+                    submit_btn = gr.Button("Generate SQL", variant="primary")
+
+                # SQL output display
+                sql_output = gr.Code(
+                    language="sql", label="Generated SQL Query", interactive=False
+                )
+
+        # Event handlers
+        extract_btn.click(
+            fn=upload_and_extract_schema,
+            inputs=[file_input],
+            outputs=[schema_output, db_path],
+        )
+
+        submit_btn.click(
+            fn=stream_sql,
+            inputs=[question_input, schema_output, db_path, chatbot],
+            outputs=[chatbot, sql_output],
+            api_name="generate",
+            queue=True,
+        )
+
+        # Also trigger on enter key
+        question_input.submit(
+            fn=stream_sql,
+            inputs=[question_input, schema_output, db_path, chatbot],
+            outputs=[chatbot, sql_output],
+            api_name="generate_on_submit",
+            queue=True,
+        )
+
+    return app
+
+
+# Launch the app
+if __name__ == "__main__":
+    app = create_app()
+    app.launch(share=True)