Update app.py

app.py

@@ -1,56 +1,33 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from datasets import load_dataset
-import json
-import os
 
 # Load the Spider dataset
 spider_dataset = load_dataset("spider", split='train')  # Load a subset of the dataset
 
-# Load the database schemas
-db_schemas = {}
-database_dir = 'path/to/database/folder'
-for filename in os.listdir(database_dir):
-    if filename.endswith('.json'):
-        with open(os.path.join(database_dir, filename), 'r') as file:
-            db_schema = json.load(file)
-            db_schemas[db_schema['db_id']] = db_schema
-
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
 model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
 
-def post_process_sql_query(sql_query, db_id):
-    # Modify the SQL query to match the dataset's schema
-    if db_id in db_schemas:
-        db_schema = db_schemas[db_id]
-        for table_name in db_schema['table_names_original']:
-            if "TABLE" in sql_query:
-                sql_query = sql_query.replace("TABLE", table_name)
-                break  # Assuming only one table is referenced in the query
-        for column_name in db_schema['column_names_original']:
-            if "COLUMN" in sql_query:
-                sql_query = sql_query.replace("COLUMN", column_name[1], 1)
-    return sql_query
+def generate_sql_from_spider_query(index):
+    # Get the natural language query from the Spider dataset
+    if index < 0 or index >= len(spider_dataset):
+        return "Invalid index. Please enter a value between 0 and {}.".format(len(spider_dataset) - 1)
 
-def generate_sql_from_user_input(query, db_id):
-    # Generate SQL for the user's query
-    input_text = "translate English to SQL: " + query
+    nl_query = spider_dataset[index]['question']
+    input_text = "translate English to SQL: " + nl_query
     inputs = tokenizer(input_text, return_tensors="pt", padding=True)
     outputs = model.generate(**inputs, max_length=512)
     sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Post-process the SQL query to match the dataset's schema
-    sql_query = post_process_sql_query(sql_query, db_id)
-    return sql_query
+    return {"Natural Language Query": nl_query, "Generated SQL Query": sql_query}
 
 # Create a Gradio interface
 interface = gr.Interface(
-    fn=lambda query, db_id: generate_sql_from_user_input(query, db_id),
-    inputs=[gr.Textbox(label="Enter your natural language query"), gr.Dropdown(label="Select Database ID", choices=list(db_schemas.keys()))],
-    outputs=gr.Textbox(label="Generated SQL Query"),
+    fn=generate_sql_from_spider_query,
+    inputs=gr.Number(label="Enter the index of the query in the Spider dataset (0 to {})".format(len(spider_dataset) - 1)),
+    outputs=[gr.Textbox(label="Natural Language Query"), gr.Textbox(label="Generated SQL Query")],
     title="NL to SQL using Spider Dataset",
-    description="This interface generates an SQL query from your natural language input based on the Spider dataset."
+    description="This interface generates an SQL query from a natural language query in the Spider dataset."
 )
 
 # Launch the app