Update app.py

app.py

@@ -10,11 +10,10 @@ db_table_names = set()
 column_names = set()
 for item in spider_dataset:
     db_id = item['db_id']
-    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])
+    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
 tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")