Update app.py

app.py

@@ -1,57 +1,49 @@
-import gradio as gr
-import pandas as pd
-import sqlite3
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
-# Load NL2SQL model
-MODEL_NAME = "mrm8488/t5-base-finetuned-wikiSQL"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-
-def nl_to_sql(nl_query, uploaded_file):
-    try:
-        # Load CSV
-        df = pd.read_csv(uploaded_file.name)
-        table_name = "user_table"
-
-        # Save DataFrame to SQLite
-        conn = sqlite3.connect(":memory:")
-        df.to_sql(table_name, conn, index=False, if_exists='replace')
-
-        # Generate SQL
-        input_ids = tokenizer(nl_query, return_tensors="pt").input_ids
-        outputs = model.generate(input_ids)
-        sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Adjust SQL to use uploaded table name if needed
-        sql_query = sql_query.replace("table", table_name)
-
-        # Execute SQL safely
-        try:
-            result_df = pd.read_sql_query(sql_query, conn)
-            result = result_df.to_string(index=False)
-        except Exception as e:
-            result = f"SQL execution failed:\n{e}\n\nGenerated SQL:\n{sql_query}"
-
-        return sql_query, result
-
-    except Exception as e:
-        return "", f"Error processing file: {e}"
-
-# Gradio UI
-iface = gr.Interface(
-    fn=nl_to_sql,
-    inputs=[
-        gr.Textbox(label="Enter your question in natural language"),
-        gr.File(label="Upload CSV", type="file")
-    ],
-    outputs=[
-        gr.Textbox(label="Generated SQL Query"),
-        gr.Textbox(label="Query Result")
-    ],
-    title="Natural Language to SQL",
-    description="Upload a CSV dataset and ask questions in English. SQL queries and results will be generated automatically."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+import pandas as pd
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load NL2SQL model from Hugging Face (no API key needed)
+model_name = "PaulGan1/t5-small-spider-sql"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Global dataframe
+df = None
+
+def load_data(file):
+    global df
+    df = pd.read_csv(file.name)
+    return f"✅ Data uploaded successfully! Columns: {', '.join(df.columns)}"
+
+def generate_sql(question):
+    if df is None:
+        return "⚠️ Please upload a CSV file first."
+    
+    # Create prompt
+    columns = ", ".join(df.columns)
+    prompt = f"translate English to SQL: {question} | table: {columns}"
+    
+    # Generate SQL
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+    outputs = model.generate(**inputs, max_length=128)
+    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Execute SQL
+    try:
+        result = pd.read_sql_query(sql_query, con=df)
+        return f"🧠 SQL Query:\n{sql_query}\n\n📊 Result:\n{result.head()}"
+    except Exception:
+        return f"🧠 SQL Query:\n{sql_query}\n\n⚠️ Unable to execute SQL. (Demo only)"
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 Natural Language to SQL Query Generator")
+    file_input = gr.File(label="Upload your CSV file")
+    upload_output = gr.Textbox(label="Upload Status")
+    question = gr.Textbox(label="Ask your question in natural language:")
+    sql_output = gr.Textbox(label="Generated SQL Query & Output", lines=10)
+
+    file_input.change(load_data, inputs=file_input, outputs=upload_output)
+    question.submit(generate_sql, inputs=question, outputs=sql_output)
+
+demo.launch()