viewer

app.py

@@ -1,18 +1,49 @@
+import os
 import gradio as gr
 from datasets import load_dataset
 
-def view_dataset(sample_size):
-    dataset = load_dataset("locuslab/TOFU")
-    data = dataset['train'].select(range(sample_size)).to_pandas()
-    return data
-
-interface = gr.Interface(
-    fn=view_dataset,
-    inputs=gr.inputs.Slider(minimum=1, maximum=100, step=1, default=5, label="Number of Samples"),
-    outputs="dataframe",
-    title="TOFU Dataset Viewer",
-    description="Interactive viewer for the TOFU dataset"
-)
-
-if __name__ == "__main__":
-    interface.launch()
+tofu_ds = load_dataset('locuslab/TOFU')
+ds = tofu_ds['full']
+
+df = ds.to_pandas()
+cols = list(df.columns)
+df = df.reindex(columns=cols)
+
+LINES_NUMBER = 20
+
+def display_df():
+    df_images = df.head(LINES_NUMBER)
+    return df_images
+
+def display_next(dataframe, end):
+    start = int(end or len(dataframe))
+    end = int(start) + int(LINES_NUMBER)
+    global df
+    if end >= len(df) - 1:
+        start = 0
+        end = LINES_NUMBER
+        df = df.sample(frac=1)
+        print(f"Shuffle")
+    df_images = df.iloc[start:end]
+    assert len(df_images) == LINES_NUMBER
+    return df_images, end
+
+initial_dataframe = display_df()
+
+# Gradio Blocks
+with gr.Blocks() as demo:
+    gr.Markdown("<h1><center>TOFU Dataset Viewer</center></h1>")
+
+    with gr.Row():
+        num_end = gr.Number(visible=False)
+        b1 = gr.Button("Get Initial dataframe")
+        b2 = gr.Button("Next Rows")
+
+    with gr.Row():
+        out_dataframe = gr.Dataframe(initial_dataframe, wrap=True, interactive=False, datatype = ['str', 'str'])
+
+    b1.click(fn=display_df, outputs=out_dataframe, api_name="initial_dataframe")
+    b2.click(fn=display_next, inputs=[out_dataframe, num_end], outputs=[out_dataframe, num_end],
+             api_name="next_rows")
+
+demo.launch(debug=True, show_error=True)