Update app.py

app.py

@@ -9,9 +9,10 @@ spider_dataset = load_dataset("spider", split='train')  # Load a subset of the d
 db_table_names = set()
 column_names = set()
 for item in spider_dataset:
-    for table in item['db']['table_names_original']:
-        db_table_names.add((item['db_id'], table))
-    for column in item['db']['column_names_original']:
+    db_id = item['db_id']
+    for table in item['table_names']:
+        db_table_names.add((db_id, table))
+    for column in item['column_names']:
         column_names.add(column[1])
 
 # Load tokenizer and model