Update app.py

app.py

@@ -10,10 +10,11 @@ db_table_names = set()
 column_names = set()
 for item in spider_dataset:
     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])
+    for table in item['sql']['from']['table_units']:
+        if isinstance(table, list):
+            db_table_names.add((db_id, table[1]))
+    for column in item['sql']['select'][1]:
+        column_names.add(column[1][1])
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")