Update app.py

app.py

@@ -2,21 +2,20 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from datasets import load_dataset
 
-# Load the Spider dataset
-spider_dataset = load_dataset("wikisql", split='train')  # Load a subset of the dataset
+# Load the WikiSQL dataset
+wikisql_dataset = load_dataset("wikisql", split='train')  # Load a subset of the dataset
 
-# Extract schema information from the Spider dataset
+# Extract schema information from the WikiSQL dataset
 table_names = set()
 column_names = set()
-for item in spider_dataset:
-    for table in item['db_id']:
-        table_names.add(table)
-    for column in item['question']:
+for item in wikisql_dataset:
+    table_names.add(item['table']['name'])
+    for column in item['table']['header']:
         column_names.add(column)
 
 # Load tokenizer and model
-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
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
 
 def generate_sql_from_user_input(query):
     # Generate SQL for the user's query
@@ -40,9 +39,10 @@ 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 Spider Dataset",
-    description="This model generates an SQL query for your natural language input based on the Spider dataset."
+    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."
 )
+
 # Launch the app
 if __name__ == "__main__":
     interface.launch()