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()