Update app.py

app.py

@@ -14,7 +14,9 @@ import os
 import matplotlib.pyplot as plt
 from matplotlib.colors import ListedColormap
 
-def compute(table: np.array):
+def compute(df_table):
+    headers = df_table.columns
+    table = df_table.to_numpy()
     vfunc = np.vectorize(lambda s: len(str(s)))
     non_empty_row_mask = (vfunc(table).sum(1) != 0)
     table = table[non_empty_row_mask]
@@ -82,7 +84,6 @@ def compute(table: np.array):
 
 
 def upload_file(file, remove_entries=10):
-    global headers
     if file.name.endswith('.arff'):
         dataset = openml.datasets.OpenMLDataset('t', 'test', data_file=file.name)
         X_, _, categorical_indicator_, attribute_names_ = dataset.get_data(
@@ -101,8 +102,6 @@ def upload_file(file, remove_entries=10):
 
 
 def update_table(table):
-    global headers
-    table = pd.DataFrame(table)
     vfunc = np.vectorize(lambda s: len(str(s)))
     non_empty_row_mask = (vfunc(table).sum(1) != 0)
     table = table[non_empty_row_mask]
@@ -115,12 +114,9 @@ def update_table(table):
     eval_lines = empty_inds[0]
     
     table.iloc[eval_lines, y_column] = ''
-    table.columns = headers
     
     return table
 
-headers = []
-
 gr.Markdown("""This demo allows you to experiment with the **TabPFN** model for tabular data.
 
             If you remove values in the target column, TabPFN will make predictions on them after clicking on the Button. The first 10 target values were already removed for this example dataset, so TabPFN will predict the first 10 classes.
@@ -130,8 +126,8 @@ gr.Markdown("""This demo allows you to experiment with the **TabPFN** model for
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            inp_table = gr.DataFrame(type='numpy', value=upload_file(Path('iris.csv'), remove_entries=10)
-                                     , headers=[''] * 3)
+            inp_table = gr.DataFrame(type='pandas', value=upload_file(Path('iris.csv'), remove_entries=10)
+                                     , headers=[''] * 5)
             
             inp_file = gr.File(
             label='Drop either a .csv (without header, only numeric values for all but the labels) or a .arff file.')