Update app.py

app.py

@@ -2,20 +2,21 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from datasets import load_dataset
 
-# Load the WikiSQL dataset
-wikisql_dataset = load_dataset("wikisql", split='train')  # Load a subset of the dataset
+# Load the Spider dataset
+spider_dataset = load_dataset("spider", split='train')  # Load a subset of the dataset
 
 # Extract schema information from the dataset
 table_names = set()
 column_names = set()
-for item in wikisql_dataset:
-    table_names.add(item['table']['name'])
-    for column in item['table']['header']:
-        column_names.add(column)
+for item in spider_dataset:
+    for table in item['db']['table_names_original']:
+        table_names.add(table)
+    for column in item['db']['column_names_original']:
+        column_names.add(column[1])
 
 # Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
-model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")  # Update this to a model fine-tuned on Spider if available
+model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")  # Update this to a model fine-tuned on Spider if available
 
 def post_process_sql_query(sql_query):
     # Modify the SQL query to match the dataset's schema
@@ -45,8 +46,8 @@ interface = gr.Interface(
     fn=generate_sql_from_user_input,
     inputs=gr.Textbox(label="Enter your natural language query"),
     outputs=gr.Textbox(label="Generated SQL Query"),
-    title="NL to SQL with T5 using WikiSQL Dataset",
-    description="This model generates an SQL query for your natural language input based on the WikiSQL dataset."
+    title="NL to SQL with T5 using Spider Dataset",
+    description="This model generates an SQL query for your natural language input based on the Spider dataset."
 )
 
 # Launch the app