from datasets import load_dataset
import gradio as gr
import os
import random

auth_token = os.environ.get("token")
winoground = load_dataset("facebook/winoground", use_auth_token=auth_token)["test"]

def func(index):
    example = winoground[index]
    return example["image_0"], example["caption_0"], example["image_1"], example["caption_1"]

demo = gr.Blocks()

with demo:
    gr.Markdown("# Slide across the slider to see various examples from WinoGround")

    with gr.Column():
        slider = gr.Slider(minimum=0, maximum=400)
        with gr.Row():
            index = random.choice(range(0, 400))
            with gr.Column():
                image_input_1 = gr.Image(value=winoground[index]["image_0"])
                text_input_1 = gr.Textbox(value=winoground[index]["caption_0"])
            with gr.Column():
                image_input_2 = gr.Image(value=winoground[index]["image_1"])
                text_input_2 = gr.Textbox(value=winoground[index]["caption_1"])
    
    slider.change(func, inputs=[slider], outputs=[image_input_1, text_input_1, image_input_2, text_input_2])

demo.launch()