import gradio as gr
import datasets 
import random
import json

def load_from_ds():
    name_to_dataset = {
        "bbh": ("SaylorTwift/bbh", "test", ["sports_understanding",    "tracking_shuffled_objects_three_objects",    "navigate",    "snarks",    "date_understanding",    "reasoning_about_colored_objects",    "object_counting",    "logical_deduction_seven_objects",    "geometric_shapes",    "web_of_lies",    "movie_recommendation",    "logical_deduction_five_objects",    "salient_translation_error_detection",    "disambiguation_qa",    "temporal_sequences",    "hyperbaton",    "logical_deduction_three_objects",    "causal_judgement",    "formal_fallacies",    "tracking_shuffled_objects_seven_objects",    "ruin_names",    "penguins_in_a_table",    "boolean_expressions",    "tracking_shuffled_objects_five_objects"]),
        #"GPQA": "Idavidrein/gpqa",
        "ifeval": ("wis-k/instruction-following-eval", "train", None),
        "math": ("lighteval/MATH-Hard", "test", None),
        "mmlu_pro": ("TIGER-Lab/MMLU-Pro", "test", None),
        "musr": ("TAUR-Lab/MuSR", "murder_mysteries", ["default"]),
        "musr": ("TAUR-Lab/MuSR", "object_placements", ["default"]),
        "musr": ("TAUR-Lab/MuSR", "team_allocation", ["default"])
    }

    examples = {}
    for name, (path, split, subsets) in name_to_dataset.items():
        if subsets is None:
            ds = datasets.load_dataset(path, split=split)
            examples[name] = {"default": random.sample(list(ds), 10)}
        else:
            examples[name] = {}
            for subset in subsets:
                ds = datasets.load_dataset(path, subset, split=split)
                examples[name][subset] = random.sample(list(ds), 10)

examples = json.load(open("samples.json"))

with gr.Blocks() as demo:
    header = gr.Markdown("# Samples explorer")
    with gr.Row():
        name = gr.Dropdown(choices=examples.keys(), label="Dataset")
        subset = gr.Dropdown(choices=[], label="Subset")
        index = gr.Dropdown(choices=range(10), value=0, label="Choice index")

    with gr.Row():
        text_box = gr.JSON(value={})

    def change_subset_dropdown(name):
        keys = sorted(examples[name].keys())
        return gr.Dropdown(choices=keys, value=keys[0], label="Subset")

    name.change(fn=change_subset_dropdown, inputs=name, outputs=subset)

    for dropdown in [name, subset, index]:  
        dropdown.change(fn= lambda name, subset, index: examples[name][subset][index], inputs=[name, subset, index], outputs=text_box)

demo.launch()